chore: simplify WKSession by providing a rawSend method (#434)

src/webkit/wkAccessibility.ts

@@ -17,8 +17,8 @@ import * as accessibility from '../accessibility';
 import { WKTargetSession } from './wkConnection';
 import { Protocol } from './protocol';
 
-export async function getAccessibilityTree(sesssion: WKTargetSession) {
+export async function getAccessibilityTree(session: WKTargetSession) {
-  const {axNode} = await sesssion.send('Page.accessibilitySnapshot');
+  const {axNode} = await session.send('Page.accessibilitySnapshot');
   return new WKAXNode(axNode);
 }
 

src/webkit/wkConnection.ts

@@ -136,8 +136,8 @@ export class WKConnection extends platform.EventEmitter {
   }
 }
 
-export const WKTargetSessionEvents = {
+export const WKSessionEvents = {
-  Disconnected: Symbol('TargetSessionEvents.Disconnected')
+  Disconnected: Symbol('WKSessionEvents.Disconnected')
 };
 
 export class WKPageProxySession extends platform.EventEmitter {
@@ -175,12 +175,12 @@ export class WKPageProxySession extends platform.EventEmitter {
     if (object.method === 'Target.targetCreated') {
       const targetInfo = object.params.targetInfo as Protocol.Target.TargetInfo;
       const session = new WKTargetSession(this, targetInfo);
-      this._sessions.set(session._sessionId, session);
+      this._sessions.set(session.sessionId, session);
       Promise.resolve().then(() => this.emit(WKPageProxySessionEvents.TargetCreated, session, object.params.targetInfo));
     } else if (object.method === 'Target.targetDestroyed') {
       const session = this._sessions.get(object.params.targetId);
       if (session) {
-        session._onClosed();
+        session.dispose();
         this._sessions.delete(object.params.targetId);
       }
       Promise.resolve().then(() => this.emit(WKPageProxySessionEvents.TargetDestroyed, { targetId: object.params.targetId, crashed: object.params.crashed }));
@@ -192,7 +192,7 @@ export class WKPageProxySession extends platform.EventEmitter {
       if (session.isProvisional())
         session._addProvisionalMessage(message);
       else
-        session._dispatchMessageFromTarget(message);
+        session.dispatchMessage(JSON.parse(message));
     } else if (object.method === 'Target.didCommitProvisionalTarget') {
       const {oldTargetId, newTargetId} = object.params as Protocol.Target.didCommitProvisionalTargetPayload;
       Promise.resolve().then(() => this.emit(WKPageProxySessionEvents.DidCommitProvisionalTarget, { oldTargetId, newTargetId }));
@@ -202,9 +202,11 @@ export class WKPageProxySession extends platform.EventEmitter {
       const oldSession = this._sessions.get(oldTargetId);
       if (!oldSession)
         throw new Error('Unknown old target: ' + oldTargetId);
-      oldSession._swappedOut = true;
+      // TODO: make some calls like screenshot catch swapped out error and retry.
+      oldSession.errorText = 'Target was swapped out.';
+      assert(newSession.isProvisional());
       for (const message of newSession._takeProvisionalMessagesAndCommit())
-        newSession._dispatchMessageFromTarget(message);
+        newSession.dispatchMessage(JSON.parse(message));
     } else {
       Promise.resolve().then(() => this.emit(object.method, object.params));
     }
@@ -216,7 +218,7 @@ export class WKPageProxySession extends platform.EventEmitter {
 
   dispose() {
     for (const session of this._sessions.values())
-      session._onClosed();
+      session.dispose();
     this._sessions.clear();
 
     this._closePromiseCallback();
@@ -225,22 +227,55 @@ export class WKPageProxySession extends platform.EventEmitter {
 }
 
 export class WKSession extends platform.EventEmitter {
+  connection: WKConnection | null;
+  readonly sessionId: string;
+  private _rawSend: (message: any) => void;
+  errorText: string;
   readonly _callbacks = new Map<number, {resolve:(o: any) => void, reject: (e: Error) => void, error: Error, method: string}>();
+
   on: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   addListener: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   off: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   removeListener: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
   once: <T extends keyof Protocol.Events | symbol>(event: T, listener: (payload: T extends symbol ? any : Protocol.Events[T extends keyof Protocol.Events ? T : never]) => void) => this;
 
+  constructor(connection: WKConnection, sessionId: string, errorText: string, rawSend: (message: any) => void) {
+    super();
+    this.connection = connection;
+    this.sessionId = sessionId;
+    this._rawSend = rawSend;
+    this.errorText = errorText;
+  }
+
   send<T extends keyof Protocol.CommandParameters>(
     method: T,
     params?: Protocol.CommandParameters[T]
   ): Promise<Protocol.CommandReturnValues[T]> {
-    throw new Error('Not implemented');
+    if (!this.connection)
+      return Promise.reject(new Error(`Protocol error (${method}): ${this.errorText}`));
+    const id = this.connection.nextMessageId();
+    const messageObj = { id, method, params };
+    debugWrappedMessage('SEND ► ' + JSON.stringify(messageObj, null, 2));
+    const result = new Promise<Protocol.CommandReturnValues[T]>((resolve, reject) => {
+      this._callbacks.set(id, {resolve, reject, error: new Error(), method});
+    });
+    this._rawSend(messageObj);
+    return result;
   }
 
-  protected _dispatchMessage(message: string) {
+  isDisposed(): boolean {
-    const object = JSON.parse(message);
+    return !this.connection;
+  }
+
+  dispose() {
+    for (const callback of this._callbacks.values())
+      callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): ${this.errorText}`));
+    this._callbacks.clear();
+    this.connection = null;
+    Promise.resolve().then(() => this.emit(WKSessionEvents.Disconnected));
+  }
+
+  dispatchMessage(object: any) {
     debugWrappedMessage('◀ RECV ' + JSON.stringify(object, null, 2));
     if (object.id && this._callbacks.has(object.id)) {
       const callback = this._callbacks.get(object.id);
@@ -249,27 +284,27 @@ export class WKSession extends platform.EventEmitter {
         callback.reject(createProtocolError(callback.error, callback.method, object));
       else
         callback.resolve(object.result);
+    } else if (object.id) {
+      // Response might come after session has been disposed and rejected all callbacks.
+      assert(this.isDisposed());
     } else {
-      assert(!object.id);
       Promise.resolve().then(() => this.emit(object.method, object.params));
     }
   }
 }
 
 export class WKTargetSession extends WKSession {
-  _pageProxySession: WKPageProxySession;
-  private readonly _targetType: string;
-  readonly _sessionId: string;
-  _swappedOut = false;
   private _provisionalMessages?: string[];
 
   constructor(pageProxySession: WKPageProxySession, targetInfo: Protocol.Target.TargetInfo) {
-    super();
+    super(pageProxySession._connection, targetInfo.targetId, `The ${targetInfo.type} has been closed.`, (message: any) => {
-    const {targetId, type, isProvisional} = targetInfo;
+      pageProxySession.send('Target.sendMessageToTarget', {
-    this._pageProxySession = pageProxySession;
+        message: JSON.stringify(message), targetId: targetInfo.targetId
-    this._targetType = type;
+      }).catch(e => {
-    this._sessionId = targetId;
+        this.dispatchMessage({ id: message.id, error: { message: e.message } });
-    if (isProvisional)
+      });
+    });
+    if (targetInfo.isProvisional)
       this._provisionalMessages = [];
   }
 
@@ -277,40 +312,6 @@ export class WKTargetSession extends WKSession {
     return !!this._provisionalMessages;
   }
 
-  isClosed(): boolean {
-    return !this._pageProxySession;
-  }
-
-  send<T extends keyof Protocol.CommandParameters>(
-    method: T,
-    params?: Protocol.CommandParameters[T]
-  ): Promise<Protocol.CommandReturnValues[T]> {
-    if (!this._pageProxySession)
-      return Promise.reject(new Error(`Protocol error (${method}): Session closed. Most likely the ${this._targetType} has been closed.`));
-    const innerId = this._pageProxySession._connection.nextMessageId();
-    const messageObj = {
-      id: innerId,
-      method,
-      params
-    };
-    debugWrappedMessage('SEND ► ' + JSON.stringify(messageObj, null, 2));
-    // Serialize message before adding callback in case JSON throws.
-    const message = JSON.stringify(messageObj);
-    const result = new Promise<Protocol.CommandReturnValues[T]>((resolve, reject) => {
-      this._callbacks.set(innerId, {resolve, reject, error: new Error(), method});
-    });
-    this._pageProxySession.send('Target.sendMessageToTarget', {
-      message: message, targetId: this._sessionId
-    }).catch(e => {
-      // There is a possible race of the connection closure. We may have received
-      // targetDestroyed notification before response for the command, in that
-      // case it's safe to swallow the exception.
-      const callback = this._callbacks.get(innerId);
-      assert(!callback, 'Callback was not rejected when target was destroyed.');
-    });
-    return result;
-  }
-
   _addProvisionalMessage(message: string) {
     this._provisionalMessages.push(message);
   }
@@ -320,24 +321,6 @@ export class WKTargetSession extends WKSession {
     this._provisionalMessages = undefined;
     return messages;
   }
-
-  _dispatchMessageFromTarget(message: string) {
-    console.assert(!this.isProvisional());
-    this._dispatchMessage(message);
-  }
-
-  _onClosed() {
-    for (const callback of this._callbacks.values()) {
-      // TODO: make some calls like screenshot catch swapped out error and retry.
-      if (this._swappedOut)
-        callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target was swapped out.`));
-      else
-        callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Target closed.`));
-    }
-    this._callbacks.clear();
-    this._pageProxySession = null;
-    Promise.resolve().then(() => this.emit(WKTargetSessionEvents.Disconnected));
-  }
 }
 
 export function createProtocolError(error: Error, method: string, object: { error: { message: string; data: any; }; }): Error {

src/webkit/wkNetworkManager.ts

@@ -95,7 +95,7 @@ export class WKNetworkManager {
     const frame = this._page._frameManager.frame(event.frameId);
     // TODO(einbinder) this will fail if we are an XHR document request
     const isNavigationRequest = event.type === 'Document';
-    const documentId = isNavigationRequest ? this._session._sessionId + '::' + event.loaderId : undefined;
+    const documentId = isNavigationRequest ? this._session.sessionId + '::' + event.loaderId : undefined;
     const request = new InterceptableRequest(this._session, this._page._state.interceptNetwork, frame, event, redirectChain, documentId);
     this._requestIdToRequest.set(event.requestId, request);
     this._page._frameManager.requestStarted(request.request);

src/webkit/wkPage.ts

@@ -19,7 +19,7 @@ import * as frames from '../frames';
 import { debugError, helper, RegisteredListener } from '../helper';
 import * as dom from '../dom';
 import * as network from '../network';
-import { WKTargetSession, WKTargetSessionEvents, WKPageProxySession } from './wkConnection';
+import { WKTargetSession, WKSessionEvents, WKPageProxySession } from './wkConnection';
 import { Events } from '../events';
 import { WKExecutionContext, EVALUATION_SCRIPT_URL } from './wkExecutionContext';
 import { WKNetworkManager } from './wkNetworkManager';
@@ -82,7 +82,6 @@ export class WKPage implements PageDelegate {
     this._addSessionListeners();
     this._networkManager.setSession(session);
     this._workers.setSession(session);
-    this._page._clearWorkers();
     this._isolatedWorlds = new Set();
     // New bootstrap scripts may have been added during provisional load, push them
     // again to be on the safe side.
@@ -116,7 +115,7 @@ export class WKPage implements PageDelegate {
     if (this._page._state.extraHTTPHeaders !== null)
       promises.push(this._setExtraHTTPHeaders(session, this._page._state.extraHTTPHeaders));
     await Promise.all(promises).catch(e => {
-      if (session.isClosed())
+      if (session.isDisposed())
         return;
       // Swallow initialization errors due to newer target swap in,
       // since we will reinitialize again.
@@ -148,7 +147,7 @@ export class WKPage implements PageDelegate {
       helper.addEventListener(this._session, 'Console.messageAdded', event => this._onConsoleMessage(event)),
       helper.addEventListener(this._pageProxySession, 'Dialog.javascriptDialogOpening', event => this._onDialog(event)),
       helper.addEventListener(this._session, 'Page.fileChooserOpened', event => this._onFileChooserOpened(event)),
-      helper.addEventListener(this._session, WKTargetSessionEvents.Disconnected, event => this._page._didDisconnect()),
+      helper.addEventListener(this._session, WKSessionEvents.Disconnected, event => this._page._didDisconnect()),
     ];
   }
 
@@ -192,7 +191,7 @@ export class WKPage implements PageDelegate {
       }
     }
     // Append session id to avoid cross-process loaderId clash.
-    const documentId = this._session._sessionId + '::' + framePayload.loaderId;
+    const documentId = this._session.sessionId + '::' + framePayload.loaderId;
     this._page._frameManager.frameCommittedNewDocumentNavigation(framePayload.id, framePayload.url, framePayload.name || '', documentId, initial);
   }
 
@@ -377,8 +376,8 @@ export class WKPage implements PageDelegate {
   }
 
   async closePage(runBeforeUnload: boolean): Promise<void> {
-    this._session._pageProxySession.send('Target.close', {
+    this._pageProxySession.send('Target.close', {
-      targetId: this._session._sessionId,
+      targetId: this._session.sessionId,
       runBeforeUnload
     }).catch(debugError);
   }

src/webkit/wkWorkers.ts

@@ -14,15 +14,16 @@
  * limitations under the License.
  */
 
-import { assert, helper, RegisteredListener } from '../helper';
+import { helper, RegisteredListener } from '../helper';
 import { Page, Worker } from '../page';
 import { Protocol } from './protocol';
-import { rewriteError, WKSession, WKTargetSession } from './wkConnection';
+import { WKSession, WKTargetSession } from './wkConnection';
 import { WKExecutionContext } from './wkExecutionContext';
 
 export class WKWorkers {
   private _sessionListeners: RegisteredListener[] = [];
   private _page: Page;
+  private _workerSessions = new Map<string, WKSession>();
 
   constructor(page: Page) {
     this._page = page;
@@ -30,10 +31,20 @@ export class WKWorkers {
 
   setSession(session: WKTargetSession) {
     helper.removeEventListeners(this._sessionListeners);
+    this._page._clearWorkers();
+    this._workerSessions.clear();
     this._sessionListeners = [
       helper.addEventListener(session, 'Worker.workerCreated', async (event: Protocol.Worker.workerCreatedPayload) => {
         const worker = new Worker(event.url);
-        const workerSession = new WKWorkerSession(session, event.workerId);
+        const workerSession = new WKSession(session.connection, event.workerId, 'Most likely the worker has been closed.', (message: any) => {
+          session.send('Worker.sendMessageToWorker', {
+            workerId: event.workerId,
+            message: JSON.stringify(message)
+          }).catch(e => {
+            workerSession.dispatchMessage({ id: message.id, error: { message: e.message } });
+          });
+        });
+        this._workerSessions.set(event.workerId, workerSession);
         worker._createExecutionContext(new WKExecutionContext(workerSession, undefined));
         this._page._addWorker(event.workerId, worker);
         workerSession.on('Console.messageAdded', event => this._onConsoleMessage(worker, event));
@@ -49,7 +60,14 @@ export class WKWorkers {
           // Worker can go as we are initializing it.
         }
       }),
+      helper.addEventListener(session, 'Worker.dispatchMessageFromWorker', (event: Protocol.Worker.dispatchMessageFromWorkerPayload) => {
+        const workerSession = this._workerSessions.get(event.workerId);
+        workerSession.dispatchMessage(JSON.parse(event.message));
+      }),
       helper.addEventListener(session, 'Worker.workerTerminated', (event: Protocol.Worker.workerTerminatedPayload) => {
+        const workerSession = this._workerSessions.get(event.workerId);
+        workerSession.dispose();
+        this._workerSessions.delete(event.workerId);
         this._page._removeWorker(event.workerId);
       })
     ];
@@ -73,59 +91,3 @@ export class WKWorkers {
     this._page._addConsoleMessage(derivedType, handles, { url, lineNumber: lineNumber - 1, columnNumber: columnNumber - 1 }, handles.length ? undefined : text);
   }
 }
-
-export class WKWorkerSession extends WKSession {
-  private _targetSession: WKTargetSession | null;
-  private _workerId: string;
-  private _lastId = 1001;
-
-  constructor(targetSession: WKTargetSession, workerId: string) {
-    super();
-    this._targetSession = targetSession;
-    this._workerId = workerId;
-    this._targetSession.on('Worker.dispatchMessageFromWorker', event => {
-      if (event.workerId === workerId)
-        this._dispatchMessage(event.message);
-    });
-    this._targetSession.on('Worker.workerTerminated', event => {
-      if (event.workerId === workerId)
-        this._workerTerminated();
-    });
-  }
-
-  send<T extends keyof Protocol.CommandParameters>(
-    method: T,
-    params?: Protocol.CommandParameters[T]
-  ): Promise<Protocol.CommandReturnValues[T]> {
-    if (!this._targetSession)
-      return Promise.reject(new Error(`Protocol error (${method}):  Most likely the worker has been closed.`));
-    const innerId = ++this._lastId;
-    const messageObj = {
-      id: innerId,
-      method,
-      params
-    };
-    const message = JSON.stringify(messageObj);
-    const result = new Promise<Protocol.CommandReturnValues[T]>((resolve, reject) => {
-      this._callbacks.set(innerId, {resolve, reject, error: new Error(), method});
-    });
-    this._targetSession.send('Worker.sendMessageToWorker', {
-      workerId: this._workerId,
-      message: message
-    }).catch(e => {
-      // There is a possible race of the connection closure. We may have received
-      // targetDestroyed notification before response for the command, in that
-      // case it's safe to swallow the exception.
-      const callback = this._callbacks.get(innerId);
-      assert(!callback, 'Callback was not rejected when worker was terminated.');
-    });
-    return result;
-  }
-
-  _workerTerminated() {
-    for (const callback of this._callbacks.values())
-      callback.reject(rewriteError(callback.error, `Protocol error (${callback.method}): Worker terminated.`));
-    this._callbacks.clear();
-    this._targetSession = null;
-  }
-}