@livestore/webmesh 0.3.0-dev.2 → 0.3.0-dev.21

package/src/websocket-connection.ts

@@ -1,19 +1,19 @@
-import type { Scope } from '@livestore/utils/effect'
 import {
   Deferred,
   Effect,
   Either,
-  FiberHandle,
+  Exit,
   Queue,
   Schedule,
   Schema,
+  Scope,
   Stream,
   WebChannel,
   WebSocket,
 } from '@livestore/utils/effect'
-import type NodeWebSocket from 'ws'
+import type * as NodeWebSocket from 'ws'
 
-import * as MeshSchema from './mesh-schema.js'
+import * as WebmeshSchema from './mesh-schema.js'
 import type { MeshNode } from './node.js'
 
 export class WSConnectionInit extends Schema.TaggedStruct('WSConnectionInit', {
@@ -48,111 +48,112 @@ export const connectViaWebSocket = ({
   reconnect?: Schedule.Schedule<unknown> | false
 }): Effect.Effect<void, never, Scope.Scope> =>
   Effect.gen(function* () {
-    const fiberHandle = yield* FiberHandle.make()
+    const disconnected = yield* Deferred.make<void>()
 
-    const connect = Effect.gen(function* () {
-      const socket = yield* WebSocket.makeWebSocket({ url, reconnect })
+    const socket = yield* WebSocket.makeWebSocket({ url, reconnect })
 
-      // NOTE we want to use `runFork` here so this Effect is not part of the fiber that will be interrupted
-      socket.addEventListener('close', () => FiberHandle.run(fiberHandle, connect).pipe(Effect.runFork))
+    socket.addEventListener('close', () => Deferred.unsafeDone(disconnected, Exit.void))
 
-      const connection = yield* makeWebSocketConnection(socket, { _tag: 'leaf', from: node.nodeName })
+    const connection = yield* makeWebSocketConnection(socket, { _tag: 'leaf', from: node.nodeName })
 
-      yield* node.addConnection({ target: 'ws', connectionChannel: connection.webChannel, replaceIfExists: true })
+    yield* node.addConnection({ target: 'ws', connectionChannel: connection.webChannel, replaceIfExists: true })
 
-      yield* Effect.never
-    }).pipe(Effect.scoped, Effect.withSpan('@livestore/webmesh:websocket-connection:connect'))
-
-    yield* FiberHandle.run(fiberHandle, connect)
-  })
+    yield* disconnected
+  }).pipe(Effect.scoped, Effect.forever, Effect.catchTag('WebSocketError', Effect.orDie))
 
 export const makeWebSocketConnection = (
   socket: globalThis.WebSocket | NodeWebSocket.WebSocket,
   socketType: SocketType,
 ): Effect.Effect<
   {
-    webChannel: WebChannel.WebChannel<typeof MeshSchema.Packet.Type, typeof MeshSchema.Packet.Type>
+    webChannel: WebChannel.WebChannel<typeof WebmeshSchema.Packet.Type, typeof WebmeshSchema.Packet.Type>
     from: string
   },
   never,
   Scope.Scope
 > =>
-  Effect.gen(function* () {
-    socket.binaryType = 'arraybuffer'
-
-    const fromDeferred = yield* Deferred.make<string>()
-
-    yield* Stream.fromEventListener<MessageEvent>(socket as any, 'message').pipe(
-      Stream.map((msg) => Schema.decodeUnknownEither(MessageMsgPack)(new Uint8Array(msg.data))),
-      Stream.flatten(),
-      Stream.tap((msg) =>
+  Effect.scopeWithCloseable((scope) =>
+    Effect.gen(function* () {
+      socket.binaryType = 'arraybuffer'
+
+      const fromDeferred = yield* Deferred.make<string>()
+
+      const listenQueue = yield* Queue.unbounded<typeof WebmeshSchema.Packet.Type>().pipe(
+        Effect.acquireRelease(Queue.shutdown),
+      )
+
+      const schema = WebChannel.mapSchema(WebmeshSchema.Packet)
+
+      yield* Stream.fromEventListener<MessageEvent>(socket as any, 'message').pipe(
+        Stream.map((msg) => Schema.decodeUnknownEither(MessageMsgPack)(new Uint8Array(msg.data))),
+        Stream.flatten(),
+        Stream.tap((msg) =>
+          Effect.gen(function* () {
+            if (msg._tag === 'WSConnectionInit') {
+              yield* Deferred.succeed(fromDeferred, msg.from)
+            } else {
+              const decodedPayload = yield* Schema.decode(schema.listen)(msg.payload)
+              yield* Queue.offer(listenQueue, decodedPayload)
+            }
+          }),
+        ),
+        Stream.runDrain,
+        Effect.interruptible,
+        Effect.tapCauseLogPretty,
+        Effect.forkScoped,
+      )
+
+      const initHandshake = (from: string) =>
+        socket.send(Schema.encodeSync(MessageMsgPack)({ _tag: 'WSConnectionInit', from }))
+
+      if (socketType._tag === 'leaf') {
+        initHandshake(socketType.from)
+      }
+
+      const deferredResult = yield* fromDeferred
+      const from = socketType._tag === 'leaf' ? socketType.from : deferredResult
+
+      if (socketType._tag === 'relay') {
+        initHandshake(from)
+      }
+
+      const isConnectedLatch = yield* Effect.makeLatch(true)
+
+      const closedDeferred = yield* Deferred.make<void>().pipe(Effect.acquireRelease(Deferred.done(Exit.void)))
+
+      yield* Effect.eventListener<any>(
+        socket,
+        'close',
+        () =>
+          Effect.gen(function* () {
+            yield* isConnectedLatch.close
+            yield* Deferred.succeed(closedDeferred, undefined)
+          }),
+        { once: true },
+      )
+
+      const send = (message: typeof WebmeshSchema.Packet.Type) =>
         Effect.gen(function* () {
-          if (msg._tag === 'WSConnectionInit') {
-            yield* Deferred.succeed(fromDeferred, msg.from)
-          } else {
-            const decodedPayload = yield* Schema.decode(MeshSchema.Packet)(msg.payload)
-            yield* Queue.offer(listenQueue, decodedPayload)
-          }
-        }),
-      ),
-      Stream.runDrain,
-      Effect.tapCauseLogPretty,
-      Effect.forkScoped,
-    )
-
-    const listenQueue = yield* Queue.unbounded<typeof MeshSchema.Packet.Type>().pipe(
-      Effect.acquireRelease(Queue.shutdown),
-    )
-
-    const initHandshake = (from: string) =>
-      socket.send(Schema.encodeSync(MessageMsgPack)({ _tag: 'WSConnectionInit', from }))
-
-    if (socketType._tag === 'leaf') {
-      initHandshake(socketType.from)
-    }
-
-    const deferredResult = yield* fromDeferred
-    const from = socketType._tag === 'leaf' ? socketType.from : deferredResult
-
-    if (socketType._tag === 'relay') {
-      initHandshake(from)
-    }
-
-    const isConnectedLatch = yield* Effect.makeLatch(true)
-
-    const closedDeferred = yield* Deferred.make<void>()
-
-    yield* Effect.eventListener<any>(
-      socket,
-      'close',
-      () =>
-        Effect.gen(function* () {
-          yield* isConnectedLatch.close
-          yield* Deferred.succeed(closedDeferred, undefined)
-        }),
-      { once: true },
-    )
-
-    const send = (message: typeof MeshSchema.Packet.Type) =>
-      Effect.gen(function* () {
-        yield* isConnectedLatch.await
-        const payload = yield* Schema.encode(MeshSchema.Packet)(message)
-        socket.send(Schema.encodeSync(MessageMsgPack)({ _tag: 'WSConnectionPayload', payload, from }))
-      })
-
-    const listen = Stream.fromQueue(listenQueue).pipe(Stream.map(Either.right))
-
-    const webChannel = {
-      [WebChannel.WebChannelSymbol]: WebChannel.WebChannelSymbol,
-      send,
-      listen,
-      closedDeferred,
-      schema: { listen: MeshSchema.Packet, send: MeshSchema.Packet },
-      supportsTransferables: false,
-    } satisfies WebChannel.WebChannel<typeof MeshSchema.Packet.Type, typeof MeshSchema.Packet.Type>
-
-    return {
-      webChannel: webChannel as WebChannel.WebChannel<typeof MeshSchema.Packet.Type, typeof MeshSchema.Packet.Type>,
-      from,
-    }
-  })
+          yield* isConnectedLatch.await
+          const payload = yield* Schema.encode(schema.send)(message)
+          socket.send(Schema.encodeSync(MessageMsgPack)({ _tag: 'WSConnectionPayload', payload, from }))
+        })
+
+      const listen = Stream.fromQueue(listenQueue).pipe(
+        Stream.map(Either.right),
+        WebChannel.listenToDebugPing('websocket-connection'),
+      )
+
+      const webChannel = {
+        [WebChannel.WebChannelSymbol]: WebChannel.WebChannelSymbol,
+        send,
+        listen,
+        closedDeferred,
+        schema,
+        supportsTransferables: false,
+        shutdown: Scope.close(scope, Exit.void),
+      } satisfies WebChannel.WebChannel<typeof WebmeshSchema.Packet.Type, typeof WebmeshSchema.Packet.Type>
+
+      return { webChannel, from }
+    }).pipe(Effect.withSpanScoped('makeWebSocketConnection')),
+  )

package/src/websocket-server.ts

@@ -1,5 +1,6 @@
+import { UnexpectedError } from '@livestore/common'
 import type { Scope } from '@livestore/utils/effect'
-import { Effect } from '@livestore/utils/effect'
+import { Effect, FiberSet } from '@livestore/utils/effect'
 import * as WebSocket from 'ws'
 
 import { makeMeshNode } from './node.js'
@@ -13,10 +14,26 @@ export const makeWebSocketServer = ({
   Effect.gen(function* () {
     const server = new WebSocket.WebSocketServer({ noServer: true })
 
+    yield* Effect.addFinalizer(() =>
+      Effect.async<void, UnexpectedError>((cb) => {
+        server.close((cause) => {
+          if (cause) {
+            cb(Effect.fail(UnexpectedError.make({ cause })))
+          } else {
+            server.removeAllListeners()
+            server.clients.forEach((client) => client.terminate())
+            cb(Effect.succeed(undefined))
+          }
+        })
+      }).pipe(Effect.orDie),
+    )
+
     const node = yield* makeMeshNode(relayNodeName)
 
     const runtime = yield* Effect.runtime<never>()
 
+    const fiberSet = yield* FiberSet.make()
+
     // TODO handle node disconnects (i.e. remove respective connection)
     server.on('connection', (socket) => {
       Effect.gen(function* () {
@@ -29,11 +46,11 @@ export const makeWebSocketServer = ({
           Effect.gen(function* () {
             yield* node.removeConnection(from)
             yield* Effect.log(`WS Relay ${relayNodeName}: removed connection from '${from}'`)
-          }).pipe(Effect.provide(runtime), Effect.runFork),
+          }).pipe(Effect.provide(runtime), Effect.tapCauseLogPretty, Effect.runFork),
         )
 
         yield* Effect.never
-      }).pipe(Effect.scoped, Effect.tapCauseLogPretty, Effect.provide(runtime), Effect.runFork)
+      }).pipe(Effect.scoped, Effect.tapCauseLogPretty, Effect.provide(runtime), FiberSet.run(fiberSet), Effect.runFork)
     })
 
     return server