[do] Broker moved to transport

Author: jamesgpearce

src/durable-object/broker.ts

@@ -10,7 +10,6 @@ import {
   getClientId,
   getPathId,
 } from './common.ts';
-import {SyncletDurableObject} from './synclet.ts';
 
 export const createDurableObjectBrokerTransport: typeof createDurableObjectBrokerTransportDecl =
   ({path, brokerPaths, ...options}) => {
@@ -27,62 +26,48 @@ export const createDurableObjectBrokerTransport: typeof createDurableObjectBroke
 
     const sendPacket = async (_packet: string): Promise<void> => {};
 
-    const fetch = async (_request: Request): Promise<Response | undefined> => {
-      return undefined;
+    const fetch = async (
+      ctx: DurableObjectState,
+      request: Request,
+    ): Promise<Response | undefined> => {
+      const pathId = getPathId(request);
+      return ifNotNull(
+        getClientId(request),
+        (clientId) => {
+          const [webSocket, client] = objValues(new WebSocketPair());
+          ctx.acceptWebSocket(client, [clientId, pathId]);
+          return createResponse(101, webSocket);
+        },
+        createUpgradeRequiredResponse,
+      ) as Response;
     };
 
+    const webSocketMessage = async (
+      ctx: DurableObjectState,
+      ws: WebSocket,
+      message: ArrayBuffer | string,
+    ): Promise<boolean | undefined> =>
+      ifNotUndefined(ctx.getTags(ws)[0], (clientId) => {
+        const packet = message.toString();
+        const splitAt = packet.indexOf(SPACE);
+        if (splitAt !== -1) {
+          const to = slice(packet, 0, splitAt);
+          const remainder = slice(packet, splitAt + 1);
+          const forwardedPacket = clientId + SPACE + remainder;
+          if (to === ASTERISK) {
+            arrayForEach(ctx.getWebSockets(), (otherClient) =>
+              otherClient !== ws ? otherClient.send(forwardedPacket) : 0,
+            );
+          } else if (to != clientId) {
+            ctx.getWebSockets(to)[0]?.send(forwardedPacket);
+          }
+        }
+        return true;
+      });
+
     return createDurableObjectTransport(
       {connect, disconnect, sendPacket},
-      {fetch},
+      {fetch, webSocketMessage},
       options,
     );
   };
-
-export class BrokerOnlyDurableObject<
-  Env = unknown,
-> extends SyncletDurableObject<Env> {
-  constructor(ctx: DurableObjectState, env: Env) {
-    super(ctx, env);
-  }
-
-  async fetch(request: Request): Promise<Response> {
-    const pathId = getPathId(request);
-    return ifNotNull(
-      getClientId(request),
-      (clientId) => {
-        const [webSocket, client] = objValues(new WebSocketPair());
-        this.ctx.acceptWebSocket(client, [clientId, pathId]);
-        return createResponse(101, webSocket);
-      },
-      createUpgradeRequiredResponse,
-    ) as Response;
-  }
-
-  webSocketMessage(client: WebSocket, message: ArrayBuffer | string) {
-    ifNotUndefined(this.ctx.getTags(client)[0], (clientId) =>
-      this.#handleMessage(clientId, message.toString(), client),
-    );
-  }
-
-  webSocketClose(_client: WebSocket) {}
-
-  #handleMessage(id: string, packet: string, fromClient?: WebSocket) {
-    const splitAt = packet.indexOf(SPACE);
-    if (splitAt !== -1) {
-      const to = slice(packet, 0, splitAt);
-      const remainder = slice(packet, splitAt + 1);
-      const forwardedPacket = id + SPACE + remainder;
-      if (to === ASTERISK) {
-        arrayForEach(this.#getClients(), (otherClient) =>
-          otherClient !== fromClient ? otherClient.send(forwardedPacket) : 0,
-        );
-      } else if (to != id) {
-        this.#getClients(to)[0]?.send(forwardedPacket);
-      }
-    }
-  }
-
-  #getClients(tag?: string) {
-    return this.ctx.getWebSockets(tag);
-  }
-}

src/durable-object/common.ts

@@ -10,14 +10,27 @@ const PATH_REGEX = /\/([^?]*)/;
 
 export const createDurableObjectTransport = (
   {connect, disconnect, sendPacket}: TransportImplementations,
-  {fetch}: {fetch: (request: Request) => Promise<Response | undefined>},
+  {
+    fetch,
+    webSocketMessage,
+  }: {
+    fetch: (
+      ctx: DurableObjectState,
+      request: Request,
+    ) => Promise<Response | undefined>;
+    webSocketMessage?: (
+      ctx: DurableObjectState,
+      client: WebSocket,
+      message: ArrayBuffer | string,
+    ) => Promise<boolean | undefined>;
+  },
   {durableObject, ...options}: DurableObjectTransportOptions,
 ) => {
   const getDurableObject = () => durableObject;
 
   return createTransport({connect, disconnect, sendPacket}, options, {
     _brand2: 'DurableObjectTransport',
-    __: [fetch],
+    __: [fetch, webSocketMessage],
     getDurableObject,
   }) as DurableObjectTransport;
 };

src/durable-object/index.ts

@@ -1,7 +1,4 @@
-export {
-  BrokerOnlyDurableObject,
-  createDurableObjectBrokerTransport,
-} from './broker.ts';
+export {createDurableObjectBrokerTransport} from './broker.ts';
 export {
   createDurableObjectSqliteDataConnector,
   createDurableObjectSqliteMetaConnector,

src/durable-object/synclet.ts

@@ -25,8 +25,6 @@ import {
 } from './common.ts';
 import {ProtectedDurableObjectTransport} from './types.ts';
 
-const FETCH = 0;
-
 export abstract class SyncletDurableObject<
   Env = unknown,
   Depth extends number = number,
@@ -63,15 +61,30 @@ export abstract class SyncletDurableObject<
   }
 
   async fetch(request: Request): Promise<Response> {
-    for (const durableObjectTransport of this.#durableObjectTransports) {
-      const response = await durableObjectTransport.__[FETCH](request);
+    for (const {
+      __: [fetch],
+    } of this.#durableObjectTransports) {
+      const response = await fetch(this.ctx, request);
       if (!isUndefined(response)) {
         return response;
       }
     }
     return createNotImplementedResponse();
   }
 
+  async webSocketMessage(
+    ws: WebSocket,
+    message: string | ArrayBuffer,
+  ): Promise<void> {
+    for (const {
+      __: [, webSocketMessage],
+    } of this.#durableObjectTransports) {
+      if (await webSocketMessage(this.ctx, ws, message)) {
+        break;
+      }
+    }
+  }
+
   getCreateDataConnector?(): DataConnectorType;
 
   getCreateMetaConnector?(): MetaConnectorType;

src/durable-object/types.ts

@@ -1,5 +1,15 @@
 import {DurableObjectTransport} from '@synclets/@types/durable-object';
 
 export interface ProtectedDurableObjectTransport extends DurableObjectTransport {
-  __: [fetch: (request: Request) => Promise<Response | undefined>];
+  __: [
+    fetch: (
+      ctx: DurableObjectState,
+      request: Request,
+    ) => Promise<Response | undefined>,
+    webSocketMessage: (
+      ctx: DurableObjectState,
+      client: WebSocket,
+      message: ArrayBuffer | string,
+    ) => Promise<boolean | undefined>,
+  ];
 }

test/durable-object/broker-only.test.ts

@@ -63,8 +63,8 @@ test('accept WebSocket upgrade requests', async () => {
 });
 
 test('2 clients communicate', async () => {
-  const a = await createClients(2);
-  const [[webSocket1, webSocket2], [received1, received2]] = a;
+  const [[webSocket1, webSocket2], [received1, received2]] =
+    await createClients(2);
 
   expect(await api('getClientCount')).toEqual(2);
 
@@ -95,6 +95,9 @@ test('2 clients communicate', async () => {
     [client1Id, 'from1To*'],
     [client1Id, 'from1To2'],
   ]);
+
+  webSocket1.close();
+  webSocket2.close();
 });
 
 test('3 clients communicate', async () => {
@@ -103,6 +106,8 @@ test('3 clients communicate', async () => {
     [received1, received2, received3],
   ] = await createClients(3);
 
+  expect(await api('getClientCount')).toEqual(3);
+
   webSocket1.send('* from1To*');
   await pause(10);
   webSocket2.send('* from2To*');
@@ -155,4 +160,8 @@ test('3 clients communicate', async () => {
     [client1Id, 'from1To3'],
     [client2Id, 'from2To3'],
   ]);
+
+  webSocket1.close();
+  webSocket2.close();
+  webSocket3.close();
 });

test/durable-object/servers.ts

@@ -1,5 +1,4 @@
 import {
-  BrokerOnlyDurableObject,
   createDurableObjectBrokerTransport,
   createDurableObjectSqliteDataConnector,
   createDurableObjectSqliteMetaConnector,
@@ -21,45 +20,30 @@ export class TestSyncletDurableObject extends SyncletDurableObject {
   async api(method: string, ...args: any[]): Promise<any> {
     return await api(this, method, ...args);
   }
-}
-
-export class TestBrokerOnlyDurableObject extends BrokerOnlyDurableObject {
-  async api(method: string, ...args: any[]): Promise<any> {
-    return await api(this, method, ...args);
-  }
 
   getClientCount() {
     return this.ctx.getWebSockets().length;
   }
 }
 
-export class TestBrokerOnlyDurableObject2 extends SyncletDurableObject {
-  getCreateComponents() {
-    return {
-      transport: createDurableObjectBrokerTransport({durableObject: this}),
-    };
-  }
-
-  async api(method: string, ...args: any[]): Promise<any> {
-    return await api(this, method, ...args);
-  }
-
-  getClientCount() {
-    return this.ctx.getWebSockets().length;
+export class TestBrokerOnlyDurableObject extends TestSyncletDurableObject {
+  getCreateTransport() {
+    return createDurableObjectBrokerTransport({durableObject: this});
   }
 }
 
 export class TestConnectorsOnlyDurableObject extends TestSyncletDurableObject {
-  getCreateComponents() {
-    return {
-      dataConnector: createDurableObjectSqliteDataConnector({
-        depth: 1,
-        sqlStorage: this.ctx.storage.sql,
-      }),
-      metaConnector: createDurableObjectSqliteMetaConnector({
-        depth: 1,
-        sqlStorage: this.ctx.storage.sql,
-      }),
-    };
+  getCreateDataConnector() {
+    return createDurableObjectSqliteDataConnector({
+      depth: 1,
+      sqlStorage: this.ctx.storage.sql,
+    });
+  }
+
+  getCreateMetaConnector() {
+    return createDurableObjectSqliteMetaConnector({
+      depth: 1,
+      sqlStorage: this.ctx.storage.sql,
+    });
   }
 }