[Reporting] Optimize visualizations awaiter performance (#118012)

x-pack/plugins/reporting/server/browsers/chromium/driver/chromium_driver.ts

@@ -210,36 +210,16 @@ export class HeadlessChromiumDriver {
     return resp;
   }
 
-  public async waitFor(
-    {
-      fn,
-      args,
-      toEqual,
-      timeout,
-    }: {
-      fn: EvaluateFn;
-      args: SerializableOrJSHandle[];
-      toEqual: number;
-      timeout: number;
-    },
-    context: EvaluateMetaOpts,
-    logger: LevelLogger
-  ): Promise<void> {
-    const startTime = Date.now();
-
-    while (true) {
-      const result = await this.evaluate({ fn, args }, context, logger);
-      if (result === toEqual) {
-        return;
-      }
-
-      if (Date.now() - startTime > timeout) {
-        throw new Error(
-          `Timed out waiting for the items selected to equal ${toEqual}. Found: ${result}. Context: ${context.context}`
-        );
-      }
-      await new Promise((r) => setTimeout(r, WAIT_FOR_DELAY_MS));
-    }
+  public async waitFor({
+    fn,
+    args,
+    timeout,
+  }: {
+    fn: EvaluateFn;
+    args: SerializableOrJSHandle[];
+    timeout: number;
+  }): Promise<void> {
+    await this.page.waitForFunction(fn, { timeout, polling: WAIT_FOR_DELAY_MS }, ...args);
   }
 
   public async setViewport(

x-pack/plugins/reporting/server/lib/screenshots/wait_for_visualizations.ts

@@ -11,10 +11,23 @@ import { HeadlessChromiumDriver } from '../../browsers';
 import { LayoutInstance } from '../layouts';
 import { CONTEXT_WAITFORELEMENTSTOBEINDOM } from './constants';
 
-type SelectorArgs = Record<string, string>;
+interface CompletedItemsCountParameters {
+  context: string;
+  count: number;
+  renderCompleteSelector: string;
+}
 
-const getCompletedItemsCount = ({ renderCompleteSelector }: SelectorArgs) => {
-  return document.querySelectorAll(renderCompleteSelector).length;
+const getCompletedItemsCount = ({
+  context,
+  count,
+  renderCompleteSelector,
+}: CompletedItemsCountParameters) => {
+  const { length } = document.querySelectorAll(renderCompleteSelector);
+
+  // eslint-disable-next-line no-console
+  console.debug(`evaluate ${context}: waitng for ${count} elements, got ${length}.`);
+
+  return length >= count;
 };
 
 /*
@@ -40,11 +53,11 @@ export const waitForVisualizations = async (
   );
 
   try {
-    await browser.waitFor(
-      { fn: getCompletedItemsCount, args: [{ renderCompleteSelector }], toEqual, timeout },
-      { context: CONTEXT_WAITFORELEMENTSTOBEINDOM },
-      logger
-    );
+    await browser.waitFor({
+      fn: getCompletedItemsCount,
+      args: [{ renderCompleteSelector, context: CONTEXT_WAITFORELEMENTSTOBEINDOM, count: toEqual }],
+      timeout,
+    });
 
     logger.debug(`found ${toEqual} rendered elements in the DOM`);
   } catch (err) {