Remove listeners on stop

Author: airhorns

spec/MultiTabWorkerBroker.spec.ts

@@ -9,6 +9,26 @@ describe("MultiTabWorkerBroker", () => {
   let testCounter = 0;
   let getLockName: () => string;
 
+  // Helper function to wait for a condition with retry logic
+  async function waitFor<T>(
+    fn: () => T | Promise<T>,
+    predicate: (value: T) => boolean,
+    options: { timeout?: number; interval?: number; message?: string } = {}
+  ): Promise<T> {
+    const { timeout = 5000, interval = 10, message = "Condition not met within timeout" } = options;
+    const startTime = Date.now();
+
+    while (Date.now() - startTime < timeout) {
+      const value = await fn();
+      if (predicate(value)) {
+        return value;
+      }
+      await new Promise((resolve) => setTimeout(resolve, interval));
+    }
+
+    throw new Error(message);
+  }
+
   beforeEach(() => {
     // Create unique lock name for each test
     const lockName = `test-lock-${testCounter++}`;
@@ -215,8 +235,14 @@ describe("MultiTabWorkerBroker", () => {
 
       // Wait for follower to become leader and initialize its worker
       await broker2LeaderPromise;
-      // Give a bit more time for worker to be fully ready
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      // Wait for broker2 to be fully ready as leader
+      await waitFor(
+        () => broker2.isLeader,
+        (isLeader) => isLeader === true,
+        { timeout: 500, message: "Broker2 did not become leader" }
+      );
+      // Give worker additional time to be fully ready
+      await new Promise((resolve) => setTimeout(resolve, 100));
       expect(broker2.isLeader).toBe(true);
 
       // Verify new leader can communicate
@@ -226,7 +252,13 @@ describe("MultiTabWorkerBroker", () => {
         method: "echo",
         params: { promoted: true },
       } as any);
-      await new Promise((resolve) => setTimeout(resolve, 50));
+
+      // Wait for the response with retry logic
+      await waitFor(
+        () => followerMessages,
+        (msgs) => msgs.some((m) => m.id === 1 && m.result?.promoted === true),
+        { timeout: 2000, message: "Expected response from promoted leader not received" }
+      );
 
       expect(followerMessages).toContainEqual(expect.objectContaining({ id: 1, result: { promoted: true } }));
 
@@ -518,6 +550,52 @@ describe("MultiTabWorkerBroker", () => {
       await broker2.stop();
     });
 
+    it("should let a follower send messages after leader restarts within same tab", async () => {
+      const lockName = getLockName();
+
+      // Tab A: start, stop (without awaiting), and start again (React Strict Mode pattern)
+      const leader = new MultiTabWorkerBroker(lockName, makeWorker);
+      await leader.start();
+      const stopPromise = leader.stop();
+      await leader.start();
+
+      expect(leader.isLeader).toBe(true);
+
+      // Tab B: start as follower
+      const follower = new MultiTabWorkerBroker(lockName, makeWorker, { timeout: 50 });
+      await follower.start();
+      expect(follower.isLeader).toBe(false);
+
+      const followerMessages: any[] = [];
+      const followerConnection = follower.createConnection();
+      followerConnection.reader.listen((msg) => followerMessages.push(msg));
+
+      // First request from follower should succeed (currently fails with "Broker stopped" due to concurrent stop cleanup)
+      const writePromise = followerConnection.writer.write({
+        jsonrpc: "2.0",
+        id: 1,
+        method: "echo",
+        params: { tab: "follower" },
+      } as any);
+
+      try {
+        await expect(writePromise).resolves.toBeUndefined();
+
+        // Wait for response delivery
+        await waitFor(
+          () => followerMessages,
+          (msgs) => msgs.some((msg) => msg.id === 1 && msg.result?.tab === "follower"),
+          { timeout: 2000, message: "Follower did not receive response after leader restart" }
+        );
+
+        expect(followerMessages).toContainEqual(expect.objectContaining({ id: 1, result: { tab: "follower" } }));
+      } finally {
+        await follower.stop();
+        await leader.stop();
+        await Promise.race([stopPromise, new Promise((resolve) => setTimeout(resolve, 1000))]);
+      }
+    });
+
     it("should timeout when no leader available", async () => {
       const lockName = getLockName();
       // Use a shorter timeout (100ms) for faster test execution
@@ -1022,10 +1100,14 @@ describe("MultiTabWorkerBroker", () => {
 
       const leader = new MultiTabWorkerBroker(lockName, async () => {
         // Delay creating the worker to widen the window for queued requests
-        await new Promise((r) => setTimeout(r, 100));
+        await new Promise((r) => setTimeout(r, 200));
+        return await makeSlowWorker();
+      });
+      const follower = new MultiTabWorkerBroker(lockName, async () => {
+        // Also delay follower's worker creation to test queuing
+        await new Promise((r) => setTimeout(r, 200));
         return await makeSlowWorker();
       });
-      const follower = new MultiTabWorkerBroker(lockName, makeSlowWorker);
 
       const followerMsgs: any[] = [];
       const followerConn = follower.createConnection();
@@ -1042,15 +1124,28 @@ describe("MultiTabWorkerBroker", () => {
       // Force leader to release so follower acquires lock and begins booting
       const stopPromise = leader.stop();
       // Small delay to ensure follower's lock wait proceeds
-      await new Promise((r) => setTimeout(r, 10));
+      await new Promise((r) => setTimeout(r, 20));
       // Now, before new worker is ready, send a couple of requests
       const p1 = followerConn.writer.write({ jsonrpc: "2.0", id: 1, method: "echo", params: { during: 1 } } as any);
       const p2 = followerConn.writer.write({ jsonrpc: "2.0", id: 2, method: "echo", params: { during: 2 } } as any);
       await stopPromise;
 
       // Give time for follower to acquire lock and broker to create worker, then for queue to drain
       await Promise.all([p1, p2]);
-      await new Promise((r) => setTimeout(r, 200));
+
+      // Wait for follower to become leader (needs more time due to 200ms worker delay)
+      await waitFor(
+        () => follower.isLeader,
+        (isLeader) => isLeader === true,
+        { timeout: 3000, message: "Follower did not become leader" }
+      );
+
+      // Wait for both responses with retry logic (allow more time for queued messages to process)
+      await waitFor(
+        () => followerMsgs,
+        (msgs) => msgs.some((m) => m.id === 1 && m.result?.during === 1) && msgs.some((m) => m.id === 2 && m.result?.during === 2),
+        { timeout: 3000, message: "Expected responses not received during leader boot" }
+      );
 
       expect(follower.isLeader).toBe(true);
       // Both responses should arrive even though they were sent while leader was booting

src/MultiTabWorkerBroker.ts

@@ -22,14 +22,14 @@ export class MultiTabWorkerBrokerError extends Error {
  */
 export class MultiTabWorkerBroker {
   isLeader = false;
+  started = false;
 
   private readonly brokerId: string;
   private broadcastChannel: BroadcastChannel | null = null;
   private worker: Worker | null = null;
   private lockAbortController = new AbortController();
   private pendingRequests = new Map<string, { resolve: (message: Message) => void; reject: (error: Error) => void }>();
   private nextRequestId = 0;
-  private started = false;
   private timeout: number;
   private onStateChange?: (state: { isLeader: boolean }) => void;
   // Queue of worker requests received while leader is booting the worker
@@ -44,6 +44,8 @@ export class MultiTabWorkerBroker {
   private shouldDebug = false;
   // Track the active lock request promise to ensure proper cleanup
   private activeLockPromise: Promise<any> | null = null;
+  // Track an in-flight stop operation so subsequent starts wait for cleanup
+  private stopPromise: Promise<void> | null = null;
 
   constructor(
     private readonly lockName: string,
@@ -104,6 +106,9 @@ export class MultiTabWorkerBroker {
 
   /** Start the broker and attempt to acquire leadership */
   async start(): Promise<void> {
+    if (this.stopPromise) {
+      await this.stopPromise;
+    }
     if (this.started) {
       return;
     }
@@ -114,7 +119,7 @@ export class MultiTabWorkerBroker {
 
     // Set up broadcast channel for inter-tab communication
     this.broadcastChannel = new BroadcastChannel(this.lockName);
-    this.broadcastChannel.addEventListener("message", this.handleBroadcastMessage.bind(this));
+    this.broadcastChannel.addEventListener("message", this.handleBroadcastMessage);
 
     // Try to acquire the lock and become leader - wait until we know our role
     await this.tryAcquireLock();
@@ -352,7 +357,12 @@ export class MultiTabWorkerBroker {
     this.emitErrorToAllConnections(error);
   }
 
-  private handleBroadcastMessage(event: MessageEvent): void {
+  private handleBroadcastMessage = (event: MessageEvent): void => {
+    // Ignore messages if we've been stopped
+    if (!this.started) {
+      return;
+    }
+
     const brokerMessage = event.data as BrokerMessage;
 
     if (brokerMessage.type === "worker-message") {
@@ -389,7 +399,7 @@ export class MultiTabWorkerBroker {
         this.pendingRequests.delete(brokerMessage.id);
       }
     }
-  }
+  };
 
   private async sendMessage(message: Message): Promise<void> {
     const originalId = (message as any)?.id;
@@ -459,73 +469,81 @@ export class MultiTabWorkerBroker {
 
   /** Stop the broker and release all resources */
   async stop(): Promise<void> {
+    if (this.stopPromise) {
+      await this.stopPromise;
+      return;
+    }
+
     if (!this.started) {
       return;
     }
 
     this.started = false;
 
-    // Release the lock
-    if (this.lockAbortController) {
-      this.lockAbortController.abort();
-    }
+    const performStop = async () => {
+      if (this.lockAbortController) {
+        this.lockAbortController.abort();
+      }
 
-    // Wait for the lock to be fully released before proceeding
-    // This is critical to ensure a new broker can immediately acquire the same lock
-    const lockPromiseToAwait = this.activeLockPromise;
-    if (lockPromiseToAwait) {
-      try {
-        await lockPromiseToAwait;
-      } catch (error) {
-        // Ignore abort errors which are expected during stop
-        if (error && (error as any).name !== "AbortError") {
-          this.error("Error while waiting for lock release:", error);
+      const lockPromiseToAwait = this.activeLockPromise;
+      if (lockPromiseToAwait) {
+        try {
+          await lockPromiseToAwait;
+        } catch (error) {
+          if (error && (error as any).name !== "AbortError") {
+            this.error("Error while waiting for lock release:", error as Error);
+          }
         }
+        this.activeLockPromise = null;
       }
-      this.activeLockPromise = null;
-    }
 
-    // Terminate worker if we're the leader
-    if (this.worker) {
-      this.worker.terminate();
-      this.worker = null;
-    }
+      if (this.worker) {
+        this.worker.terminate();
+        this.worker = null;
+      }
 
-    // Close broadcast channel
-    if (this.broadcastChannel) {
-      this.broadcastChannel.close();
-      this.broadcastChannel = null;
-    }
+      const channel = this.broadcastChannel;
+      if (channel) {
+        channel.removeEventListener("message", this.handleBroadcastMessage);
+        channel.close();
+        this.broadcastChannel = null;
+      }
 
-    // Reject all pending requests
-    for (const [id, pending] of this.pendingRequests.entries()) {
-      pending.reject(new Error("Broker stopped"));
-    }
-    this.pendingRequests.clear();
+      for (const [id, pending] of this.pendingRequests.entries()) {
+        pending.reject(new Error("Broker stopped"));
+      }
+      this.pendingRequests.clear();
 
-    // Emit close events to all connections
-    this.emitCloseToAllConnections();
+      this.emitCloseToAllConnections();
 
-    // Dispose all connections
-    for (const { reader, writer } of this.connections.values()) {
-      reader.dispose();
-      writer.dispose();
-    }
-    this.connections.clear();
+      for (const { reader, writer } of this.connections.values()) {
+        reader.dispose();
+        writer.dispose();
+      }
+      this.connections.clear();
+
+      this.workerReadyPromise = null;
+      this.workerReadyResolver = null;
+      this.leaderRequestQueue = [];
+      this.rewrittenIdMap.clear();
+
+      const wasLeader = this.isLeader;
+      this.isLeader = false;
 
-    // Reset worker state
-    this.workerReadyPromise = null;
-    this.workerReadyResolver = null;
-    this.leaderRequestQueue = [];
-    this.rewrittenIdMap.clear();
+      if (wasLeader) {
+        this.onStateChange?.({ isLeader: false });
+      }
+    };
 
-    const wasLeader = this.isLeader;
-    this.isLeader = false;
+    this.stopPromise = (async () => {
+      try {
+        await performStop();
+      } finally {
+        this.stopPromise = null;
+      }
+    })();
 
-    // Notify state change if we were the leader
-    if (wasLeader) {
-      this.onStateChange?.({ isLeader: false });
-    }
+    await this.stopPromise;
   }
 }