From d5fef5f78f259389ce73bb6887b918bdc8176a5e Mon Sep 17 00:00:00 2001
From: Michael Lin <unadlib@gmail.com>
Date: Sun, 8 Sep 2024 18:02:45 +0800
Subject: [PATCH] feat(webworker): support onConnect for webworker (#4)

---
 examples/webworker/src/index.ts   | 17 ++++---
 examples/webworker/src/worker.ts  |  9 +++-
 src/transports/workerTransport.ts | 74 +++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/examples/webworker/src/index.ts b/examples/webworker/src/index.ts
index b914698..8e8cccb 100644
--- a/examples/webworker/src/index.ts
+++ b/examples/webworker/src/index.ts
@@ -3,7 +3,8 @@ import { Worker, Main } from './interface';
 
 class MainTransport
   extends WorkerTransport.Main<{ emit: Main }>
-  implements Worker {
+  implements Worker
+{
   async help() {
     const response = await this.emit('help', { text: 'SOS!!!' });
     return response;
@@ -21,13 +22,19 @@ class MainTransport
 
 const worker = new Worker('worker.bundle.js');
 
-(window as any).mainTransport = new MainTransport({
-  worker,
-});
-
 document.getElementById('btn')?.addEventListener('click', async () => {
   const response = await (window as any).mainTransport.help();
   const div = document.createElement('div');
   div.innerText = `${new Date()}: ${JSON.stringify(response)}`;
   document.body.appendChild(div);
 });
+
+// mock async init worker
+setTimeout(() => {
+  (window as any).mainTransport = new MainTransport({
+    worker,
+  });
+  (window as any).mainTransport.onConnect(() => {
+    console.log('connected');
+  });
+}, 1000);
diff --git a/examples/webworker/src/worker.ts b/examples/webworker/src/worker.ts
index 856b1f4..bdfa490 100644
--- a/examples/webworker/src/worker.ts
+++ b/examples/webworker/src/worker.ts
@@ -18,4 +18,11 @@ class WebWorkerTransport
   }
 }
 
-(self as any).webWorkerTransport = new WebWorkerTransport();
+// mock async init worker
+setTimeout(() => {
+  (self as any).webWorkerTransport = new WebWorkerTransport();
+
+  (self as any).webWorkerTransport.onConnect(() => {
+    console.log('connected');
+  });
+}, 2000);
diff --git a/src/transports/workerTransport.ts b/src/transports/workerTransport.ts
index 0cc7345..1d43b58 100644
--- a/src/transports/workerTransport.ts
+++ b/src/transports/workerTransport.ts
@@ -10,6 +10,10 @@ import { Transport } from '../transport';
 // workaround: `tsc --skipLibCheck`.
 declare var self: WorkerGlobalScope;
 
+type ClientCallback = () => void | Promise<void>;
+
+const connectEventName = 'worker-connect';
+
 export interface WorkerMainTransportOptions
   extends Partial<TransportOptions<TransferableWorker>> {
   /**
@@ -50,6 +54,41 @@ export abstract class WorkerMainTransport<
       listener,
       sender,
     });
+
+    this.emit({
+      // @ts-ignore
+      name: connectEventName,
+      respond: true,
+      silent: true,
+    }).then(this._handleConnectCallbacks);
+
+    // @ts-ignore
+    this.listen(connectEventName, this._handleConnectCallbacks);
+  }
+
+  private _connected = false;
+
+  private _handleConnectCallbacks = async () => {
+    if (this._connected) {
+      return;
+    }
+    this._connected = true;
+    this._onConnectCallback.forEach((callback) => {
+      callback();
+    });
+    this._onConnectCallback.clear();
+  };
+
+  private _onConnectCallback = new Set<ClientCallback>();
+
+  onConnect(callback: ClientCallback) {
+    if (this._connected) {
+      return callback();
+    }
+    this._onConnectCallback.add(callback);
+    return () => {
+      this._onConnectCallback.delete(callback);
+    };
   }
 }
 
@@ -79,6 +118,41 @@ export abstract class WorkerInternalTransport<
       listener,
       sender,
     });
+
+    this.emit({
+      // @ts-ignore
+      name: connectEventName,
+      respond: true,
+      silent: true,
+    }).then(this._handleConnectCallbacks);
+
+    // @ts-ignore
+    this.listen(connectEventName, this._handleConnectCallbacks);
+  }
+
+  private _handleConnectCallbacks = async () => {
+    if (this._connected) {
+      return;
+    }
+    this._connected = true;
+    this._onConnectCallback.forEach((callback) => {
+      callback();
+    });
+    this._onConnectCallback.clear();
+  };
+
+  private _connected = false;
+
+  private _onConnectCallback = new Set<ClientCallback>();
+
+  onConnect(callback: ClientCallback) {
+    if (this._connected) {
+      return callback();
+    }
+    this._onConnectCallback.add(callback);
+    return () => {
+      this._onConnectCallback.delete(callback);
+    };
   }
 }