chore: introduce progress.runAtomic()

packages/playwright-core/src/server/frames.ts

@@ -27,7 +27,7 @@ import * as network from './network';
 import { Page } from './page';
 import { ProgressController } from './progress';
 import * as types from './types';
-import { LongStandingScope, asLocator, assert, constructURLBasedOnBaseURL, makeWaitForNextTask, monotonicTime, renderTitleForCall } from '../utils';
+import { LongStandingScope, asLocator, assert, constructURLBasedOnBaseURL, makeWaitForNextTask, renderTitleForCall } from '../utils';
 import { isSessionClosedError } from './protocolError';
 import { debugLogger } from './utils/debugLogger';
 import { eventsHelper } from './utils/eventsHelper';
@@ -1391,54 +1391,41 @@ export class Frame extends SdkObject {
   }
 
   private async _expectImpl(metadata: CallMetadata, selector: string, options: FrameExpectParams): Promise<{ matches: boolean, received?: any, log?: string[], timedOut?: boolean }> {
+    const controller = new ProgressController(metadata, this);
     const lastIntermediateResult: { received?: any, isSet: boolean } = { isSet: false };
-    try {
-      let timeout = options.timeout;
-      const start = timeout > 0 ? monotonicTime() : 0;
-
+    return await controller.run(async progress => {
       // Step 1: perform locator handlers checkpoint with a specified timeout.
-      await (new ProgressController(metadata, this)).run(async progress => {
-        progress.log(`${renderTitleForCall(metadata)}${timeout ? ` with timeout ${timeout}ms` : ''}`);
-        progress.log(`waiting for ${this._asLocator(selector)}`);
-        await this._page.performActionPreChecks(progress);
-      }, timeout);
+      progress.log(`${renderTitleForCall(metadata)}${options.timeout ? ` with timeout ${options.timeout}ms` : ''}`);
+      progress.log(`waiting for ${this._asLocator(selector)}`);
+      await this._page.performActionPreChecks(progress);
 
       // Step 2: perform one-shot expect check without a timeout.
       // Supports the case of `expect(locator).toBeVisible({ timeout: 1 })`
       // that should succeed when the locator is already visible.
       try {
-        const resultOneShot = await (new ProgressController(metadata, this)).run(async progress => {
-          return await this._expectInternal(progress, selector, options, lastIntermediateResult);
-        });
+        const resultOneShot = await progress.runAtomic(() => this._expectInternal(progress, selector, options, lastIntermediateResult));
         if (resultOneShot.matches !== options.isNot)
           return resultOneShot;
       } catch (e) {
         if (js.isJavaScriptErrorInEvaluate(e) || isInvalidSelectorError(e))
           throw e;
         // Ignore any other errors from one-shot, we'll handle them during retries.
       }
-      if (timeout > 0) {
-        const elapsed = monotonicTime() - start;
-        timeout -= elapsed;
-      }
-      if (timeout < 0)
-        return { matches: options.isNot, log: compressCallLog(metadata.log), timedOut: true, received: lastIntermediateResult.received };
+      progress.throwIfAborted();
 
       // Step 3: auto-retry expect with increasing timeouts. Bounded by the total remaining time.
-      return await (new ProgressController(metadata, this)).run(async progress => {
-        return await this.retryWithProgressAndTimeouts(progress, [100, 250, 500, 1000], async continuePolling => {
-          await this._page.performActionPreChecks(progress);
-          const { matches, received } = await this._expectInternal(progress, selector, options, lastIntermediateResult);
-          if (matches === options.isNot) {
-            // Keep waiting in these cases:
-            // expect(locator).conditionThatDoesNotMatch
-            // expect(locator).not.conditionThatDoesMatch
-            return continuePolling;
-          }
-          return { matches, received };
-        });
-      }, timeout);
-    } catch (e) {
+      return await this.retryWithProgressAndTimeouts(progress, [100, 250, 500, 1000], async continuePolling => {
+        await this._page.performActionPreChecks(progress);
+        const { matches, received } = await this._expectInternal(progress, selector, options, lastIntermediateResult);
+        if (matches === options.isNot) {
+          // Keep waiting in these cases:
+          // expect(locator).conditionThatDoesNotMatch
+          // expect(locator).not.conditionThatDoesMatch
+          return continuePolling;
+        }
+        return { matches, received };
+      });
+    }, options.timeout).catch(e => {
       // Q: Why not throw upon isSessionClosedError(e) as in other places?
       // A: We want user to receive a friendly message containing the last intermediate result.
       if (js.isJavaScriptErrorInEvaluate(e) || isInvalidSelectorError(e))
@@ -1449,7 +1436,7 @@ export class Frame extends SdkObject {
       if (e instanceof TimeoutError)
         result.timedOut = true;
       return result;
-    }
+    });
   }
 
   private async _expectInternal(progress: Progress, selector: string, options: FrameExpectParams, lastIntermediateResult: { received?: any, isSet: boolean }) {

packages/playwright-core/src/server/progress.ts

@@ -27,6 +27,7 @@ export interface Progress {
   isRunning(): boolean;
   cleanupWhenAborted(cleanup: () => any): void;
   throwIfAborted(): void;
+  runAtomic<T>(task: () => Promise<T>): Promise<T>;
   metadata: CallMetadata;
 }
 
@@ -69,6 +70,16 @@ export class ProgressController {
     this._state = 'running';
     this.sdkObject.attribution.context?._activeProgressControllers.add(this);
 
+    let timer: NodeJS.Timeout | undefined;
+    const timeoutError = new TimeoutError(`Timeout ${this._timeout}ms exceeded.`);
+    const stopTimer = () => {
+      clearTimeout(timer);
+      timer = undefined;
+    };
+    const startTimer = () => {
+      timer = setTimeout(() => this._forceAbortPromise.reject(timeoutError), Math.max(0, progress.timeUntilDeadline()));
+    };
+
     const progress: Progress = {
       log: message => {
         if (this._state === 'running')
@@ -88,12 +99,19 @@ export class ProgressController {
         if (this._state === 'aborted')
           throw new AbortedError();
       },
-      metadata: this.metadata
+      metadata: this.metadata,
+      runAtomic: async <T>(task: () => Promise<T>): Promise<T> => {
+        stopTimer();
+        try {
+          return await task();
+        } finally {
+          startTimer();
+        }
+      },
     };
 
-    const timeoutError = new TimeoutError(`Timeout ${this._timeout}ms exceeded.`);
-    const timer = setTimeout(() => this._forceAbortPromise.reject(timeoutError), progress.timeUntilDeadline());
     try {
+      startTimer();
       const promise = task(progress);
       const result = await Promise.race([promise, this._forceAbortPromise]);
       this._state = 'finished';
@@ -104,7 +122,7 @@ export class ProgressController {
       throw e;
     } finally {
       this.sdkObject.attribution.context?._activeProgressControllers.delete(this);
-      clearTimeout(timer);
+      stopTimer();
     }
   }
 }