chore: make polling in page cancelable from node (#2399)

- unifies polling timeouts with everything else,
  based on the client time instead of the server time;
- prepares polling tasks for cancellation token
  behavior.

Unfortunately, RerunnableTask had to be rewritten almost
entirely.

src/dom.ts

@@ -113,6 +113,11 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return utility.doEvaluateInternal(true /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await utility.injectedScript(), node: this }, arg);
   }
 
+  async _evaluateHandleInUtility<R, Arg>(pageFunction: types.FuncOn<{ injected: InjectedScript, node: T }, Arg, R>, arg: Arg): Promise<js.JSHandle<R>> {
+    const utility = await this._context.frame._utilityContext();
+    return utility.doEvaluateInternal(false /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await utility.injectedScript(), node: this }, arg);
+  }
+
   async ownerFrame(): Promise<frames.Frame | null> {
     const frameId = await this._page._delegate.getOwnerFrame(this);
     if (!frameId)
@@ -349,7 +354,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     }
     return await this._page._frameManager.waitForSignalsCreatedBy<string[]>(async () => {
       const injectedResult = await this._evaluateInUtility(({ injected, node }, selectOptions) => injected.selectOptions(node, selectOptions), selectOptions);
-      return handleInjectedResult(injectedResult, '');
+      return handleInjectedResult(injectedResult);
     }, deadline, options);
   }
 
@@ -359,7 +364,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     const deadline = this._page._timeoutSettings.computeDeadline(options);
     await this._page._frameManager.waitForSignalsCreatedBy(async () => {
       const injectedResult = await this._evaluateInUtility(({ injected, node }, value) => injected.fill(node, value), value);
-      const needsInput = handleInjectedResult(injectedResult, '');
+      const needsInput = handleInjectedResult(injectedResult);
       if (needsInput) {
         if (value)
           await this._page.keyboard.insertText(value);
@@ -372,7 +377,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   async selectText(): Promise<void> {
     this._page._log(inputLog, `elementHandle.selectText()`);
     const injectedResult = await this._evaluateInUtility(({ injected, node }) => injected.selectText(node), {});
-    handleInjectedResult(injectedResult, '');
+    handleInjectedResult(injectedResult);
   }
 
   async setInputFiles(files: string | types.FilePayload | string[] | types.FilePayload[], options?: types.NavigatingActionWaitOptions) {
@@ -386,7 +391,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
       const input = node as Node as HTMLInputElement;
       return { status: 'success', value: input.multiple };
     }, {});
-    const multiple = handleInjectedResult(injectedResult, '');
+    const multiple = handleInjectedResult(injectedResult);
     let ff: string[] | types.FilePayload[];
     if (!Array.isArray(files))
       ff = [ files ] as string[] | types.FilePayload[];
@@ -414,7 +419,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   async focus() {
     this._page._log(inputLog, `elementHandle.focus()`);
     const injectedResult = await this._evaluateInUtility(({ injected, node }) => injected.focusNode(node), {});
-    handleInjectedResult(injectedResult, '');
+    handleInjectedResult(injectedResult);
   }
 
   async type(text: string, options?: { delay?: number } & types.NavigatingActionWaitOptions) {
@@ -492,13 +497,17 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   async _waitForDisplayedAtStablePositionAndEnabled(deadline: number): Promise<void> {
     this._page._log(inputLog, 'waiting for element to be displayed, enabled and not moving...');
     const rafCount =  this._page._delegate.rafCountForStablePosition();
-    const stablePromise = this._evaluateInUtility(({ injected, node }, { rafCount, timeout }) => {
+    const poll = await this._evaluateHandleInUtility(({ injected, node }, { rafCount }) => {
-      return injected.waitForDisplayedAtStablePositionAndEnabled(node, rafCount, timeout);
+      return injected.waitForDisplayedAtStablePositionAndEnabled(node, rafCount);
-    }, { rafCount, timeout: helper.timeUntilDeadline(deadline) });
+    }, { rafCount });
-    const timeoutMessage = 'element to be displayed and not moving';
+    try {
-    const injectedResult = await helper.waitWithDeadline(stablePromise, timeoutMessage, deadline, 'pw:input');
+      const stablePromise = poll.evaluate(poll => poll.result);
-    handleInjectedResult(injectedResult, timeoutMessage);
+      const injectedResult = await helper.waitWithDeadline(stablePromise, 'element to be displayed and not moving', deadline, 'pw:input');
-    this._page._log(inputLog, '...element is displayed, enabled and does not move');
+      handleInjectedResult(injectedResult);
+    } finally {
+      poll.evaluate(poll => poll.cancel()).catch(e => {}).then(() => poll.dispose());
+    }
+    this._page._log(inputLog, '...element is displayed and does not move');
   }
 
   async _checkHitTargetAt(point: types.Point): Promise<boolean> {
@@ -514,7 +523,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     const injectedResult = await this._evaluateInUtility(({ injected, node }, { point }) => {
       return injected.checkHitTargetAt(node, point);
     }, { point });
-    return handleInjectedResult(injectedResult, '');
+    return handleInjectedResult(injectedResult);
   }
 }
 
@@ -526,11 +535,9 @@ export function toFileTransferPayload(files: types.FilePayload[]): types.FileTra
   }));
 }
 
-function handleInjectedResult<T = undefined>(injectedResult: types.InjectedScriptResult<T>, timeoutMessage: string): T {
+function handleInjectedResult<T = undefined>(injectedResult: types.InjectedScriptResult<T>): T {
   if (injectedResult.status === 'notconnected')
     throw new NotConnectedError();
-  if (injectedResult.status === 'timeout')
-    throw new TimeoutError(`waiting for ${timeoutMessage} failed: timeout exceeded. Re-run with the DEBUG=pw:input env variable to see the debug log.`);
   if (injectedResult.status === 'error')
     throw new Error(injectedResult.error);
   return injectedResult.value as T;

src/frames.ts

@@ -36,7 +36,7 @@ type ContextData = {
   contextPromise: Promise<dom.FrameExecutionContext>;
   contextResolveCallback: (c: dom.FrameExecutionContext) => void;
   context: dom.FrameExecutionContext | null;
-  rerunnableTasks: Set<RerunnableTask>;
+  rerunnableTasks: Set<RerunnableTask<any>>;
 };
 
 export type GotoOptions = types.NavigateOptions & {
@@ -451,7 +451,7 @@ export class Frame {
       throw new Error(`Unsupported waitFor option "${state}"`);
 
     const deadline = this._page._timeoutSettings.computeDeadline(options);
-    const { world, task } = selectors._waitForSelectorTask(selector, state, deadline);
+    const { world, task } = selectors._waitForSelectorTask(selector, state);
     const result = await this._scheduleRerunnableTask(task, world, deadline, `selector "${selector}"${state === 'attached' ? '' : ' to be ' + state}`);
     if (!result.asElement()) {
       result.dispose();
@@ -469,7 +469,7 @@ export class Frame {
 
   async dispatchEvent(selector: string, type: string, eventInit?: Object, options?: types.TimeoutOptions): Promise<void> {
     const deadline = this._page._timeoutSettings.computeDeadline(options);
-    const task = selectors._dispatchEventTask(selector, type, eventInit || {}, deadline);
+    const task = selectors._dispatchEventTask(selector, type, eventInit || {});
     const result = await this._scheduleRerunnableTask(task, 'main', deadline, `selector "${selector}"`);
     result.dispose();
   }
@@ -700,7 +700,7 @@ export class Frame {
     this._page._log(dom.inputLog, `(page|frame).${actionName}("${selector}")`);
     while (!helper.isPastDeadline(deadline)) {
       try {
-        const { world, task } = selectors._waitForSelectorTask(selector, 'attached', deadline);
+        const { world, task } = selectors._waitForSelectorTask(selector, 'attached');
         this._page._log(dom.inputLog, `waiting for the selector "${selector}"`);
         const handle = await this._scheduleRerunnableTask(task, world, deadline, `selector "${selector}"`);
         this._page._log(dom.inputLog, `...got element for the selector`);
@@ -812,11 +812,11 @@ export class Frame {
       throw new Error('Unknown polling options: ' + polling);
     const predicateBody = helper.isString(pageFunction) ? 'return (' + pageFunction + ')' : 'return (' + pageFunction + ')(arg)';
 
-    const task = async (context: dom.FrameExecutionContext) => context.evaluateHandleInternal(({ injected, predicateBody, polling, timeout, arg }) => {
+    const task = async (context: dom.FrameExecutionContext) => context.evaluateHandleInternal(({ injected, predicateBody, polling, arg }) => {
-      const innerPredicate = new Function('arg', predicateBody);
+      const innerPredicate = new Function('arg', predicateBody) as (arg: any) => R;
-      return injected.poll(polling, timeout, () => innerPredicate(arg));
+      return injected.poll(polling, () => innerPredicate(arg));
-    }, { injected: await context.injectedScript(), predicateBody, polling, timeout: helper.timeUntilDeadline(deadline), arg });
+    }, { injected: await context.injectedScript(), predicateBody, polling, arg });
-    return this._scheduleRerunnableTask(task, 'main', deadline) as any as types.SmartHandle<R>;
+    return this._scheduleRerunnableTask(task, 'main', deadline);
   }
 
   async title(): Promise<string> {
@@ -836,10 +836,9 @@ export class Frame {
     this._parentFrame = null;
   }
 
-  private _scheduleRerunnableTask(task: Task, contextType: ContextType, deadline: number, title?: string): Promise<js.JSHandle> {
+  private _scheduleRerunnableTask<T>(task: SchedulableTask<T>, contextType: ContextType, deadline: number, title?: string): Promise<types.SmartHandle<T>> {
     const data = this._contextData.get(contextType)!;
     const rerunnableTask = new RerunnableTask(data, task, deadline, title);
-    data.rerunnableTasks.add(rerunnableTask);
     if (data.context)
       rerunnableTask.rerun(data.context);
     return rerunnableTask.promise;
@@ -890,84 +889,77 @@ export class Frame {
   }
 }
 
-type Task = (context: dom.FrameExecutionContext) => Promise<js.JSHandle>;
+export type SchedulableTask<T> = (context: dom.FrameExecutionContext) => Promise<js.JSHandle<types.CancelablePoll<T>>>;
 
-class RerunnableTask {
+class RerunnableTask<T> {
-  readonly promise: Promise<js.JSHandle>;
+  readonly promise: Promise<types.SmartHandle<T>>;
-  private _contextData: ContextData;
+  terminate: (reason: Error) => void = () => {};
-  private _task: Task;
+  private _task: SchedulableTask<T>;
-  private _runCount: number;
+  private _resolve: (result: types.SmartHandle<T>) => void = () => {};
-  private _resolve: (result: js.JSHandle) => void = () => {};
   private _reject: (reason: Error) => void = () => {};
-  private _timeoutTimer?: NodeJS.Timer;
+  private _terminatedPromise: Promise<Error>;
-  private _terminated = false;
 
-  constructor(data: ContextData, task: Task, deadline: number, title?: string) {
+  constructor(data: ContextData, task: SchedulableTask<T>, deadline: number, title?: string) {
-    this._contextData = data;
     this._task = task;
-    this._runCount = 0;
+    data.rerunnableTasks.add(this);
-    this.promise = new Promise<js.JSHandle>((resolve, reject) => {
+
-      this._resolve = resolve;
-      this._reject = reject;
-    });
     // Since page navigation requires us to re-install the pageScript, we should track
     // timeout on our end.
     const timeoutError = new TimeoutError(`waiting for ${title || 'function'} failed: timeout exceeded. Re-run with the DEBUG=pw:input env variable to see the debug log.`);
-    this._timeoutTimer = setTimeout(() => this.terminate(timeoutError), helper.timeUntilDeadline(deadline));
+    let timeoutTimer: NodeJS.Timer | undefined;
-  }
+    this._terminatedPromise = new Promise(resolve => {
+      timeoutTimer = setTimeout(() => resolve(timeoutError), helper.timeUntilDeadline(deadline));
+      this.terminate = resolve;
+    });
 
-  terminate(error: Error) {
+    // This promise is either resolved with the task result, or rejected with a meaningful
-    this._terminated = true;
+    // evaluation error.
-    this._reject(error);
+    const resultPromise = new Promise<types.SmartHandle<T>>((resolve, reject) => {
-    this._doCleanup();
+      this._resolve = resolve;
+      this._reject = reject;
+    });
+    const failPromise = this._terminatedPromise.then(error => Promise.reject(error));
+
+    this.promise = Promise.race([resultPromise, failPromise]).finally(() => {
+      if (timeoutTimer)
+        clearTimeout(timeoutTimer);
+      data.rerunnableTasks.delete(this);
+    });
   }
 
   async rerun(context: dom.FrameExecutionContext) {
-    const runCount = ++this._runCount;
+    let poll: js.JSHandle<types.CancelablePoll<T>> | null = null;
-    let success: js.JSHandle | null = null;
+
-    let error = null;
+    // On timeout or error, cancel current poll.
+    const cancelPoll = () => {
+      if (!poll)
+        return;
+      const copy = poll;
+      poll = null;
+      copy.evaluate(p => p.cancel()).catch(e => {}).then(() => copy.dispose());
+    };
+    this._terminatedPromise.then(cancelPoll);
+
     try {
-      success = await this._task(context);
+      poll = await this._task(context);
+      const result = await poll.evaluateHandle(poll => poll.result);
+      cancelPoll();
+      this._resolve(result);
     } catch (e) {
-      error = e;
+      cancelPoll();
+
+      // When the page is navigated, the promise is rejected.
+      // We will try again in the new execution context.
+      if (e.message.includes('Execution context was destroyed'))
+        return;
+
+      // We could have tried to evaluate in a context which was already
+      // destroyed.
+      if (e.message.includes('Cannot find context with specified id'))
+        return;
+
+      this._reject(e);
     }
-
-    if (this._terminated || runCount !== this._runCount) {
-      if (success)
-        success.dispose();
-      return;
-    }
-
-    // Ignore timeouts in pageScript - we track timeouts ourselves.
-    // If execution context has been already destroyed, `context.evaluate` will
-    // throw an error - ignore this predicate run altogether.
-    if (!error && await context.evaluateInternal(s => !s, success).catch(e => true)) {
-      success!.dispose();
-      return;
-    }
-
-    // When the page is navigated, the promise is rejected.
-    // We will try again in the new execution context.
-    if (error && error.message.includes('Execution context was destroyed'))
-      return;
-
-    // We could have tried to evaluate in a context which was already
-    // destroyed.
-    if (error && error.message.includes('Cannot find context with specified id'))
-      return;
-
-    if (error)
-      this._reject(error);
-    else
-      this._resolve(success!);
-
-    this._doCleanup();
-  }
-
-  _doCleanup() {
-    if (this._timeoutTimer)
-      clearTimeout(this._timeoutTimer);
-    this._contextData.rerunnableTasks.delete(this);
   }
 }
 

src/injected/injectedScript.ts

@@ -21,7 +21,15 @@ import { SelectorEngine, SelectorRoot } from './selectorEngine';
 import { createTextSelector } from './textSelectorEngine';
 import { XPathEngine } from './xpathSelectorEngine';
 
-type Predicate<T> = () => T;
+type Falsy = false | 0 | '' | undefined | null;
+type Predicate<T> = () => T | Falsy;
+type InjectedScriptProgress = {
+  canceled: boolean;
+};
+
+function asCancelablePoll<T>(result: T): types.CancelablePoll<T> {
+  return { result: Promise.resolve(result), cancel: () => {} };
+}
 
 export default class InjectedScript {
   readonly engines: Map<string, SelectorEngine>;
@@ -103,19 +111,13 @@ export default class InjectedScript {
     return rect.width > 0 && rect.height > 0;
   }
 
-  private _pollRaf<T>(predicate: Predicate<T>, timeout: number): Promise<T | undefined> {
+  private _pollRaf<T>(progress: InjectedScriptProgress, predicate: Predicate<T>): Promise<T> {
-    let timedOut = false;
+    let fulfill: (result: T) => void;
-    if (timeout)
+    const result = new Promise<T>(x => fulfill = x);
-      setTimeout(() => timedOut = true, timeout);
-
-    let fulfill: (result?: any) => void;
-    const result = new Promise<T | undefined>(x => fulfill = x);
 
     const onRaf = () => {
-      if (timedOut) {
+      if (progress.canceled)
-        fulfill();
         return;
-      }
       const success = predicate();
       if (success)
         fulfill(success);
@@ -127,18 +129,12 @@ export default class InjectedScript {
     return result;
   }
 
-  private _pollInterval<T>(pollInterval: number, predicate: Predicate<T>, timeout: number): Promise<T | undefined> {
+  private _pollInterval<T>(progress: InjectedScriptProgress, pollInterval: number, predicate: Predicate<T>): Promise<T> {
-    let timedOut = false;
+    let fulfill: (result: T) => void;
-    if (timeout)
+    const result = new Promise<T>(x => fulfill = x);
-      setTimeout(() => timedOut = true, timeout);
-
-    let fulfill: (result?: any) => void;
-    const result = new Promise<T | undefined>(x => fulfill = x);
     const onTimeout = () => {
-      if (timedOut) {
+      if (progress.canceled)
-        fulfill();
         return;
-      }
       const success = predicate();
       if (success)
         fulfill(success);
@@ -150,10 +146,11 @@ export default class InjectedScript {
     return result;
   }
 
-  poll<T>(polling: 'raf' | number, timeout: number, predicate: Predicate<T>): Promise<T | undefined> {
+  poll<T>(polling: 'raf' | number, predicate: Predicate<T>): types.CancelablePoll<T> {
-    if (polling === 'raf')
+    const progress = { canceled: false };
-      return this._pollRaf(predicate, timeout);
+    const cancel = () => { progress.canceled = true; };
-    return this._pollInterval(polling, predicate, timeout);
+    const result = polling === 'raf' ? this._pollRaf(progress, predicate) : this._pollInterval(progress, polling, predicate);
+    return { result, cancel };
   }
 
   getElementBorderWidth(node: Node): { left: number; top: number; } {
@@ -330,25 +327,25 @@ export default class InjectedScript {
     input.dispatchEvent(new Event('change', { 'bubbles': true }));
   }
 
-  async waitForDisplayedAtStablePositionAndEnabled(node: Node, rafCount: number, timeout: number): Promise<types.InjectedScriptResult> {
+  waitForDisplayedAtStablePositionAndEnabled(node: Node, rafCount: number): types.CancelablePoll<types.InjectedScriptResult> {
     if (!node.isConnected)
-      return { status: 'notconnected' };
+      return asCancelablePoll({ status: 'notconnected' });
     const element = node.nodeType === Node.ELEMENT_NODE ? (node as Element) : node.parentElement;
     if (!element)
-      return { status: 'notconnected' };
+      return asCancelablePoll({ status: 'notconnected' });
 
     let lastRect: types.Rect | undefined;
     let counter = 0;
     let samePositionCounter = 0;
     let lastTime = 0;
-    const result = await this.poll('raf', timeout, (): 'notconnected' | boolean => {
+    return this.poll('raf', (): types.InjectedScriptResult | false => {
       // First raf happens in the same animation frame as evaluation, so it does not produce
       // any client rect difference compared to synchronous call. We skip the synchronous call
       // and only force layout during actual rafs as a small optimisation.
       if (++counter === 1)
         return false;
       if (!node.isConnected)
-        return 'notconnected';
+        return { status: 'notconnected' };
 
       // Drop frames that are shorter than 16ms - WebKit Win bug.
       const time = performance.now();
@@ -373,9 +370,8 @@ export default class InjectedScript {
       const elementOrButton = element.closest('button, [role=button]') || element;
       const isDisabled = ['BUTTON', 'INPUT', 'SELECT'].includes(elementOrButton.nodeName) && elementOrButton.hasAttribute('disabled');
 
-      return isDisplayedAndStable && isVisible && !isDisabled;
+      return isDisplayedAndStable && isVisible && !isDisabled ? { status: 'success' } : false;
     });
-    return { status: result === 'notconnected' ? 'notconnected' : (result ? 'success' : 'timeout') };
   }
 
   checkHitTargetAt(node: Node, point: types.Point): types.InjectedScriptResult<boolean> {

src/selectors.ts

@@ -111,12 +111,12 @@ export class Selectors {
     return result;
   }
 
-  _waitForSelectorTask(selector: string, state: 'attached' | 'detached' | 'visible' | 'hidden', deadline: number): { world: 'main' | 'utility', task: (context: dom.FrameExecutionContext) => Promise<js.JSHandle> } {
+  _waitForSelectorTask(selector: string, state: 'attached' | 'detached' | 'visible' | 'hidden'): { world: 'main' | 'utility', task: frames.SchedulableTask<Element | boolean> } {
     const parsed = this._parseSelector(selector);
     const task = async (context: dom.FrameExecutionContext) => {
       const injectedScript = await context.injectedScript();
-      return injectedScript.evaluateHandle((injected, { parsed, state, timeout }) => {
+      return injectedScript.evaluateHandle((injected, { parsed, state }) => {
-        return injected.poll('raf', timeout, () => {
+        return injected.poll('raf', () => {
           const element = injected.querySelector(parsed, document);
           switch (state) {
             case 'attached':
@@ -129,23 +129,23 @@ export class Selectors {
               return !element || !injected.isVisible(element);
           }
         });
-      }, { parsed, state, timeout: helper.timeUntilDeadline(deadline) });
+      }, { parsed, state });
     };
     return { world: this._needsMainContext(parsed) ? 'main' : 'utility', task };
   }
 
-  _dispatchEventTask(selector: string, type: string, eventInit: Object, deadline: number): (context: dom.FrameExecutionContext) => Promise<js.JSHandle> {
+  _dispatchEventTask(selector: string, type: string, eventInit: Object): frames.SchedulableTask<Element> {
     const parsed = this._parseSelector(selector);
     const task = async (context: dom.FrameExecutionContext) => {
       const injectedScript = await context.injectedScript();
-      return injectedScript.evaluateHandle((injected, { parsed, type, eventInit, timeout }) => {
+      return injectedScript.evaluateHandle((injected, { parsed, type, eventInit }) => {
-        return injected.poll('raf', timeout, () => {
+        return injected.poll('raf', () => {
           const element = injected.querySelector(parsed, document);
           if (element)
             injected.dispatchEvent(element, type, eventInit);
           return element || false;
         });
-      }, { parsed, type, eventInit, timeout: helper.timeUntilDeadline(deadline) });
+      }, { parsed, type, eventInit });
     };
     return task;
   }

src/types.ts

@@ -167,5 +167,9 @@ export type ParsedSelector = {
 export type InjectedScriptResult<T = undefined> =
   (T extends undefined ? { status: 'success', value?: T} : { status: 'success', value: T }) |
   { status: 'notconnected' } |
-  { status: 'timeout' } |
   { status: 'error', error: string };
+
+export type CancelablePoll<T> = {
+  result: Promise<T>,
+  cancel: () => void,
+};