fix: Unsubscribe from WebSocket stream when iterator ended

neet · neet · commit b7732ed404b1 · 2024-10-15T12:45:54.000+09:00
diff --git a/src/adapters/action/dispatcher-ws.spec.ts b/src/adapters/action/dispatcher-ws.spec.ts
@@ -1,4 +1,4 @@
-import { MastoUnexpectedError } from "../errors";
+import { MastoUnexpectedError, MastoWebSocketError } from "../errors";
 import { createLogger } from "../logger";
 import { SerializerNativeImpl } from "../serializers";
 import {
@@ -29,7 +29,7 @@ describe("DispatcherWs", () => {
     }).toThrow(MastoUnexpectedError);
   });
 
-  it("can be disposed", () => {
+  it("can be disposed", async () => {
     const connector = new WebSocketConnectorImpl({
       constructorParameters: ["wss://example.com"],
     });
@@ -41,6 +41,8 @@ describe("DispatcherWs", () => {
     );
 
     dispatcher[Symbol.dispose]();
-    expect(connector.canAcquire()).toBe(false);
+    await expect(() => connector.acquire()).rejects.toThrow(
+      MastoWebSocketError,
+    );
   });
 });
diff --git a/src/adapters/ws/web-socket-connector.ts b/src/adapters/ws/web-socket-connector.ts
@@ -33,11 +33,11 @@ export class WebSocketConnectorImpl implements WebSocketConnector {
     });
   }
 
-  canAcquire(): boolean {
-    return !this.closed;
-  }
-
   async acquire(): Promise<WebSocket> {
+    if (this.closed) {
+      throw new MastoWebSocketError("WebSocket closed");
+    }
+
     this.init();
 
     if (this.ws != undefined) {
@@ -49,6 +49,12 @@ export class WebSocketConnectorImpl implements WebSocketConnector {
     return await promiseWithResolvers.promise;
   }
 
+  async *[Symbol.asyncIterator](): AsyncIterableIterator<WebSocket> {
+    while (!this.closed) {
+      yield await this.acquire();
+    }
+  }
+
   close(): void {
     this.closed = true;
     this.ws?.close();
diff --git a/src/adapters/ws/web-socket-subscription.ts b/src/adapters/ws/web-socket-subscription.ts
@@ -25,31 +25,35 @@ export class WebSocketSubscription implements mastodon.streaming.Subscription {
   async *values(): AsyncIterableIterator<mastodon.streaming.Event> {
     this.logger?.log("info", "Subscribing to stream", this.stream);
 
-    while (this.connector.canAcquire()) {
-      this.connection = await this.connector.acquire();
+    try {
+      for await (const connection of this.connector) {
+        this.connection = connection;
 
-      const data = this.serializer.serialize("json", {
-        type: "subscribe",
-        stream: this.stream,
-        ...this.params,
-      });
+        const data = this.serializer.serialize("json", {
+          type: "subscribe",
+          stream: this.stream,
+          ...this.params,
+        });
 
-      this.logger?.log("debug", "↑ WEBSOCKET", data);
-      this.connection.send(data);
-      this.counter.increment(this.stream, this.params);
+        this.logger?.log("debug", "↑ WEBSOCKET", data);
+        this.connection.send(data);
+        this.counter.increment(this.stream, this.params);
 
-      const messages = toAsyncIterable(this.connection);
+        const messages = toAsyncIterable(this.connection);
 
-      for await (const message of messages) {
-        const event = await this.parseMessage(message.data as string);
+        for await (const message of messages) {
+          const event = await this.parseMessage(message.data as string);
 
-        if (!this.test(event)) {
-          continue;
-        }
+          if (!this.test(event)) {
+            continue;
+          }
 
-        this.logger?.log("debug", "↓ WEBSOCKET", event);
-        yield event;
+          this.logger?.log("debug", "↓ WEBSOCKET", event);
+          yield event;
+        }
       }
+    } finally {
+      this.unsubscribe();
     }
   }
 
diff --git a/src/interfaces/ws.ts b/src/interfaces/ws.ts
@@ -1,9 +1,8 @@
 import { type WebSocket } from "isomorphic-ws";
 
-export interface WebSocketConnector {
+export interface WebSocketConnector extends AsyncIterable<WebSocket> {
   acquire(): Promise<WebSocket>;
   close(): void;
-  canAcquire(): boolean;
 }
 
 export interface WebSocketSubscriptionCounter {

Original file line number	Diff line number	Diff line change
`@@ -1,9 +1,8 @@`
`1`	`1`	`import { type WebSocket } from "isomorphic-ws";`
`2`	`2`
`3`		`-export interface WebSocketConnector {`
	`3`	`+export interface WebSocketConnector extends AsyncIterable<WebSocket> {`
`4`	`4`	`acquire(): Promise<WebSocket>;`
`5`	`5`	`close(): void;`
`6`		`- canAcquire(): boolean;`
`7`	`6`	`}`
`8`	`7`
`9`	`8`	`export interface WebSocketSubscriptionCounter {`