feat(rpc): merge ChannelOwner and ConnectionScope (#2911)

Author: dgozman

src/rpc/client/browser.ts

@@ -19,7 +19,6 @@ import { BrowserChannel, BrowserInitializer } from '../channels';
 import { BrowserContext } from './browserContext';
 import { Page } from './page';
 import { ChannelOwner } from './channelOwner';
-import { ConnectionScope } from './connection';
 import { Events } from '../../events';
 import { CDPSession } from './cdpSession';
 
@@ -37,13 +36,13 @@ export class Browser extends ChannelOwner<BrowserChannel, BrowserInitializer> {
     return browser ? Browser.from(browser) : null;
   }
 
-  constructor(scope: ConnectionScope, guid: string, initializer: BrowserInitializer) {
-    super(scope, guid, initializer, true);
+  constructor(parent: ChannelOwner, guid: string, initializer: BrowserInitializer) {
+    super(parent, guid, initializer, true);
     this._channel.on('close', () => {
       this._isConnected = false;
       this.emit(Events.Browser.Disconnected);
       this._isClosedOrClosing = true;
-      this._scope.dispose();
+      this._dispose();
     });
     this._closedPromise = new Promise(f => this.once(Events.Browser.Disconnected, f));
   }

src/rpc/client/browserContext.ts

@@ -23,7 +23,6 @@ import { BrowserContextChannel, BrowserContextInitializer } from '../channels';
 import { ChannelOwner } from './channelOwner';
 import { helper } from '../../helper';
 import { Browser } from './browser';
-import { ConnectionScope } from './connection';
 import { Events } from '../../events';
 import { TimeoutSettings } from '../../timeoutSettings';
 import { CDPSession } from './cdpSession';
@@ -51,8 +50,8 @@ export class BrowserContext extends ChannelOwner<BrowserContextChannel, BrowserC
     return context ? BrowserContext.from(context) : null;
   }
 
-  constructor(scope: ConnectionScope, guid: string, initializer: BrowserContextInitializer) {
-    super(scope, guid, initializer, true);
+  constructor(parent: ChannelOwner, guid: string, initializer: BrowserContextInitializer) {
+    super(parent, guid, initializer, true);
     initializer.pages.forEach(p => {
       const page = Page.from(p);
       this._pages.add(page);
@@ -225,7 +224,7 @@ export class BrowserContext extends ChannelOwner<BrowserContextChannel, BrowserC
     }
     this._pendingWaitForEvents.clear();
     this.emit(Events.BrowserContext.Close);
-    this._scope.dispose();
+    this._dispose();
   }
 
   async close(): Promise<void> {

src/rpc/client/browserServer.ts

@@ -16,7 +16,6 @@
 
 import { ChildProcess } from 'child_process';
 import { BrowserServerChannel, BrowserServerInitializer } from '../channels';
-import { ConnectionScope } from './connection';
 import { ChannelOwner } from './channelOwner';
 import { Events } from '../../events';
 
@@ -25,8 +24,8 @@ export class BrowserServer extends ChannelOwner<BrowserServerChannel, BrowserSer
     return (server as any)._object;
   }
 
-  constructor(scope: ConnectionScope, guid: string, initializer: BrowserServerInitializer) {
-    super(scope, guid, initializer);
+  constructor(parent: ChannelOwner, guid: string, initializer: BrowserServerInitializer) {
+    super(parent, guid, initializer);
     this._channel.on('close', () => this.emit(Events.BrowserServer.Close));
   }
 

src/rpc/client/browserType.ts

@@ -19,7 +19,6 @@ import { BrowserTypeChannel, BrowserTypeInitializer } from '../channels';
 import { Browser } from './browser';
 import { BrowserContext } from './browserContext';
 import { ChannelOwner } from './channelOwner';
-import { ConnectionScope } from './connection';
 import { BrowserServer } from './browserServer';
 
 export class BrowserType extends ChannelOwner<BrowserTypeChannel, BrowserTypeInitializer> {
@@ -28,8 +27,8 @@ export class BrowserType extends ChannelOwner<BrowserTypeChannel, BrowserTypeIni
     return (browserTyep as any)._object;
   }
 
-  constructor(scope: ConnectionScope, guid: string, initializer: BrowserTypeInitializer) {
-    super(scope, guid, initializer);
+  constructor(parent: ChannelOwner, guid: string, initializer: BrowserTypeInitializer) {
+    super(parent, guid, initializer);
   }
 
   executablePath(): string {

src/rpc/client/cdpSession.ts

@@ -15,7 +15,6 @@
  */
 
 import { CDPSessionChannel, CDPSessionInitializer } from '../channels';
-import { ConnectionScope } from './connection';
 import { ChannelOwner } from './channelOwner';
 import { Protocol } from '../../chromium/protocol';
 
@@ -30,11 +29,11 @@ export class CDPSession extends ChannelOwner<CDPSessionChannel, CDPSessionInitia
   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(scope: ConnectionScope, guid: string, initializer: CDPSessionInitializer) {
-    super(scope, guid, initializer, true);
+  constructor(parent: ChannelOwner, guid: string, initializer: CDPSessionInitializer) {
+    super(parent, guid, initializer, true);
 
     this._channel.on('event', ({ method, params }) => this.emit(method, params));
-    this._channel.on('disconnected', () => this._scope.dispose());
+    this._channel.on('disconnected', () => this._dispose());
 
     this.on = super.on;
     this.addListener = super.addListener;

src/rpc/client/channelOwner.ts

@@ -16,18 +16,33 @@
 
 import { EventEmitter } from 'events';
 import { Channel } from '../channels';
-import { ConnectionScope } from './connection';
+import { Connection } from './connection';
+import { assert } from '../../helper';
 
-export abstract class ChannelOwner<T extends Channel, Initializer> extends EventEmitter {
+export abstract class ChannelOwner<T extends Channel = Channel, Initializer = {}> extends EventEmitter {
+  private _connection: Connection;
+  private _isScope: boolean;
+  // Parent is always "isScope".
+  private _parent: ChannelOwner | undefined;
+  // Only "isScope" channel owners have registered objects inside.
+  private _objects = new Map<string, ChannelOwner>();
+
+  readonly _guid: string;
   readonly _channel: T;
   readonly _initializer: Initializer;
-  readonly _scope: ConnectionScope;
-  readonly guid: string;
 
-  constructor(scope: ConnectionScope, guid: string, initializer: Initializer, isScope?: boolean) {
+  constructor(parent: ChannelOwner | Connection, guid: string, initializer: Initializer, isScope?: boolean) {
     super();
-    this.guid = guid;
-    this._scope = isScope ? scope.createChild(guid) : scope;
+
+    this._connection = parent instanceof Connection ? parent : parent._connection;
+    this._guid = guid;
+    this._isScope = !!isScope;
+    this._parent = parent instanceof Connection ? undefined : parent;
+
+    this._connection._objects.set(guid, this);
+    if (this._parent)
+      this._parent._objects.set(guid, this);
+
     const base = new EventEmitter();
     this._channel = new Proxy(base, {
       get: (obj: any, prop) => {
@@ -45,10 +60,35 @@ export abstract class ChannelOwner<T extends Channel, Initializer> extends Event
           return obj.addListener;
         if (prop === 'removeEventListener')
           return obj.removeListener;
-        return (params: any) => scope.sendMessageToServer({ guid, method: String(prop), params });
+        return (params: any) => this._connection.sendMessageToServer({ guid, method: String(prop), params });
       },
     });
     (this._channel as any)._object = this;
     this._initializer = initializer;
   }
+
+  _dispose() {
+    assert(this._isScope);
+
+    // Clean up from parent and connection.
+    if (this._parent)
+      this._parent._objects.delete(this._guid);
+    this._connection._objects.delete(this._guid);
+
+    // Dispose all children.
+    for (const [guid, object] of [...this._objects]) {
+      if (object._isScope)
+        object._dispose();
+      else
+        this._connection._objects.delete(guid);
+    }
+    this._objects.clear();
+  }
+
+  _debugScopeState(): any {
+    return {
+      _guid: this._guid,
+      objects: this._isScope ? Array.from(this._objects.values()).map(o => o._debugScopeState()) : undefined,
+    };
+  }
 }