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

package/src/websocket-edge.ts

@@ -0,0 +1,159 @@
+import {
+  Deferred,
+  Effect,
+  Either,
+  Exit,
+  Queue,
+  Schedule,
+  Schema,
+  Scope,
+  Stream,
+  WebChannel,
+  WebSocket,
+} from '@livestore/utils/effect'
+import type * as NodeWebSocket from 'ws'
+
+import * as WebmeshSchema from './mesh-schema.js'
+import type { MeshNode } from './node.js'
+
+export class WSEdgeInit extends Schema.TaggedStruct('WSEdgeInit', {
+  from: Schema.String,
+}) {}
+
+export class WSEdgePayload extends Schema.TaggedStruct('WSEdgePayload', {
+  from: Schema.String,
+  payload: Schema.Any,
+}) {}
+
+export class WSEdgeMessage extends Schema.Union(WSEdgeInit, WSEdgePayload) {}
+
+export const MessageMsgPack = Schema.MsgPack(WSEdgeMessage)
+
+export type SocketType =
+  | {
+      _tag: 'leaf'
+      from: string
+    }
+  | {
+      _tag: 'relay'
+    }
+
+export const connectViaWebSocket = ({
+  node,
+  url,
+  reconnect = Schedule.exponential(100),
+}: {
+  node: MeshNode
+  url: string
+  reconnect?: Schedule.Schedule<unknown> | false
+}): Effect.Effect<void, never, Scope.Scope> =>
+  Effect.gen(function* () {
+    const disconnected = yield* Deferred.make<void>()
+
+    const socket = yield* WebSocket.makeWebSocket({ url, reconnect })
+
+    socket.addEventListener('close', () => Deferred.unsafeDone(disconnected, Exit.void))
+
+    const edgeChannel = yield* makeWebSocketEdge(socket, { _tag: 'leaf', from: node.nodeName })
+
+    yield* node.addEdge({ target: 'ws', edgeChannel: edgeChannel.webChannel, replaceIfExists: true })
+
+    yield* disconnected
+  }).pipe(Effect.scoped, Effect.forever, Effect.catchTag('WebSocketError', Effect.orDie))
+
+export const makeWebSocketEdge = (
+  socket: globalThis.WebSocket | NodeWebSocket.WebSocket,
+  socketType: SocketType,
+): Effect.Effect<
+  {
+    webChannel: WebChannel.WebChannel<typeof WebmeshSchema.Packet.Type, typeof WebmeshSchema.Packet.Type>
+    from: string
+  },
+  never,
+  Scope.Scope
+> =>
+  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 === 'WSEdgeInit') {
+              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: 'WSEdgeInit', 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* () {
+          yield* isConnectedLatch.await
+          const payload = yield* Schema.encode(schema.send)(message)
+          socket.send(Schema.encodeSync(MessageMsgPack)({ _tag: 'WSEdgePayload', payload, from }))
+        })
+
+      const listen = Stream.fromQueue(listenQueue).pipe(
+        Stream.map(Either.right),
+        WebChannel.listenToDebugPing('websocket-edge'),
+      )
+
+      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('makeWebSocketEdge')),
+  )

package/src/websocket-server.ts

@@ -1,9 +1,10 @@
+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'
-import { makeWebSocketConnection } from './websocket-connection.js'
+import { makeWebSocketEdge } from './websocket-edge.js'
 
 export const makeWebSocketServer = ({
   relayNodeName,
@@ -13,27 +14,43 @@ 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* () {
-        const { webChannel, from } = yield* makeWebSocketConnection(socket, { _tag: 'relay' })
+        const { webChannel, from } = yield* makeWebSocketEdge(socket, { _tag: 'relay' })
 
-        yield* node.addConnection({ target: from, connectionChannel: webChannel, replaceIfExists: true })
-        yield* Effect.log(`WS Relay ${relayNodeName}: added connection from '${from}'`)
+        yield* node.addEdge({ target: from, edgeChannel: webChannel, replaceIfExists: true })
+        yield* Effect.log(`WS Relay ${relayNodeName}: added edge from '${from}'`)
 
         socket.addEventListener('close', () =>
           Effect.gen(function* () {
-            yield* node.removeConnection(from)
-            yield* Effect.log(`WS Relay ${relayNodeName}: removed connection from '${from}'`)
-          }).pipe(Effect.provide(runtime), Effect.runFork),
+            yield* node.removeEdge(from)
+            yield* Effect.log(`WS Relay ${relayNodeName}: removed edge from '${from}'`)
+          }).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

package/src/websocket-connection.ts

@@ -1,158 +0,0 @@
-import type { Scope } from '@livestore/utils/effect'
-import {
-  Deferred,
-  Effect,
-  Either,
-  FiberHandle,
-  Queue,
-  Schedule,
-  Schema,
-  Stream,
-  WebChannel,
-  WebSocket,
-} from '@livestore/utils/effect'
-import type NodeWebSocket from 'ws'
-
-import * as MeshSchema from './mesh-schema.js'
-import type { MeshNode } from './node.js'
-
-export class WSConnectionInit extends Schema.TaggedStruct('WSConnectionInit', {
-  from: Schema.String,
-}) {}
-
-export class WSConnectionPayload extends Schema.TaggedStruct('WSConnectionPayload', {
-  from: Schema.String,
-  payload: Schema.Any,
-}) {}
-
-export class WSConnectionMessage extends Schema.Union(WSConnectionInit, WSConnectionPayload) {}
-
-export const MessageMsgPack = Schema.MsgPack(WSConnectionMessage)
-
-export type SocketType =
-  | {
-      _tag: 'leaf'
-      from: string
-    }
-  | {
-      _tag: 'relay'
-    }
-
-export const connectViaWebSocket = ({
-  node,
-  url,
-  reconnect = Schedule.exponential(100),
-}: {
-  node: MeshNode
-  url: string
-  reconnect?: Schedule.Schedule<unknown> | false
-}): Effect.Effect<void, never, Scope.Scope> =>
-  Effect.gen(function* () {
-    const fiberHandle = yield* FiberHandle.make()
-
-    const connect = Effect.gen(function* () {
-      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))
-
-      const connection = yield* makeWebSocketConnection(socket, { _tag: 'leaf', from: node.nodeName })
-
-      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)
-  })
-
-export const makeWebSocketConnection = (
-  socket: globalThis.WebSocket | NodeWebSocket.WebSocket,
-  socketType: SocketType,
-): Effect.Effect<
-  {
-    webChannel: WebChannel.WebChannel<typeof MeshSchema.Packet.Type, typeof MeshSchema.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.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,
-    }
-  })