@livestore/sync-cf 0.3.0-dev.16 → 0.3.0-dev.18

package/src/cf-worker/durable-object.ts

@@ -36,6 +36,8 @@ const WebSocketAttachmentSchema = Schema.parseJson(
   }),
 )
 
+export const PULL_CHUNK_SIZE = 100
+
 /**
  * Needs to be bumped when the storage format changes (e.g. mutationLogTable schema changes)
  *
@@ -45,7 +47,9 @@ export const PERSISTENCE_FORMAT_VERSION = 3
 
 export type MakeDurableObjectClassOptions = {
   onPush?: (message: WSMessage.PushReq) => Effect.Effect<void> | Promise<void>
+  onPushRes?: (message: WSMessage.PushAck | WSMessage.Error) => Effect.Effect<void> | Promise<void>
   onPull?: (message: WSMessage.PullReq) => Effect.Effect<void> | Promise<void>
+  onPullRes?: (message: WSMessage.PullRes | WSMessage.Error) => Effect.Effect<void> | Promise<void>
 }
 
 export type MakeDurableObjectClass = (options?: MakeDurableObjectClassOptions) => {
@@ -54,6 +58,9 @@ export type MakeDurableObjectClass = (options?: MakeDurableObjectClassOptions) =
 
 export const makeDurableObject: MakeDurableObjectClass = (options) => {
   return class WebSocketServerBase extends DurableObject<Env> {
+    /** Needed to prevent concurrent pushes */
+    private pushSemaphore = Effect.makeSemaphore(1).pipe(Effect.runSync)
+
     constructor(ctx: DurableObjectState, env: Env) {
       super(ctx, env)
     }
@@ -88,21 +95,21 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
         })
       }).pipe(Effect.tapCauseLogPretty, Effect.runPromise)
 
-    webSocketMessage = (ws: WebSocketClient, message: ArrayBuffer | string) =>
-      Effect.gen(this, function* () {
-        const decodedMessageRes = decodeIncomingMessage(message)
+    webSocketMessage = (ws: WebSocketClient, message: ArrayBuffer | string) => {
+      const decodedMessageRes = decodeIncomingMessage(message)
 
-        if (decodedMessageRes._tag === 'Left') {
-          console.error('Invalid message received', decodedMessageRes.left)
-          return
-        }
+      if (decodedMessageRes._tag === 'Left') {
+        console.error('Invalid message received', decodedMessageRes.left)
+        return
+      }
+
+      const decodedMessage = decodedMessageRes.right
+      const requestId = decodedMessage.requestId
 
+      return Effect.gen(this, function* () {
         const { storeId } = yield* Schema.decode(WebSocketAttachmentSchema)(ws.deserializeAttachment())
         const storage = makeStorage(this.ctx, this.env, storeId)
 
-        const decodedMessage = decodedMessageRes.right
-        const requestId = decodedMessage.requestId
-
         try {
           switch (decodedMessage._tag) {
             // TODO allow pulling concurrently to not block incoming push requests
@@ -111,8 +118,15 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
                 yield* Effect.tryAll(() => options.onPull!(decodedMessage))
               }
 
+              const respond = (message: WSMessage.PullRes) =>
+                Effect.gen(function* () {
+                  if (options?.onPullRes) {
+                    yield* Effect.tryAll(() => options.onPullRes!(message))
+                  }
+                  ws.send(encodeOutgoingMessage(message))
+                })
+
               const cursor = decodedMessage.cursor
-              const CHUNK_SIZE = 100
 
               // TODO use streaming
               const remainingEvents = yield* storage.getEvents(cursor)
@@ -121,32 +135,41 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
               const batches =
                 remainingEvents.length === 0
                   ? [[]]
-                  : Array.from({ length: Math.ceil(remainingEvents.length / CHUNK_SIZE) }, (_, i) =>
-                      remainingEvents.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE),
+                  : Array.from({ length: Math.ceil(remainingEvents.length / PULL_CHUNK_SIZE) }, (_, i) =>
+                      remainingEvents.slice(i * PULL_CHUNK_SIZE, (i + 1) * PULL_CHUNK_SIZE),
                     )
 
               for (const [index, batch] of batches.entries()) {
-                const remaining = Math.max(0, remainingEvents.length - (index + 1) * CHUNK_SIZE)
-                ws.send(encodeOutgoingMessage(WSMessage.PullRes.make({ batch, remaining })))
+                const remaining = Math.max(0, remainingEvents.length - (index + 1) * PULL_CHUNK_SIZE)
+                yield* respond(WSMessage.PullRes.make({ batch, remaining, requestId: { context: 'pull', requestId } }))
               }
 
               break
             }
             case 'WSMessage.PushReq': {
-              if (options?.onPush) {
-                yield* Effect.tryAll(() => options.onPush!(decodedMessage))
-              }
+              const respond = (message: WSMessage.PushAck | WSMessage.Error) =>
+                Effect.gen(function* () {
+                  if (options?.onPushRes) {
+                    yield* Effect.tryAll(() => options.onPushRes!(message))
+                  }
+                  ws.send(encodeOutgoingMessage(message))
+                })
 
               if (decodedMessage.batch.length === 0) {
-                ws.send(encodeOutgoingMessage(WSMessage.PushAck.make({ requestId })))
+                yield* respond(WSMessage.PushAck.make({ requestId }))
                 return
               }
 
+              yield* this.pushSemaphore.take(1)
+
+              if (options?.onPush) {
+                yield* Effect.tryAll(() => options.onPush!(decodedMessage))
+              }
+
               // TODO check whether we could use the Durable Object storage for this to speed up the lookup
               const expectedParentId = yield* storage.getHead
 
               // TODO handle clientId unique conflict
-
               // Validate the batch
               const firstEvent = decodedMessage.batch[0]!
               if (firstEvent.parentId !== expectedParentId) {
@@ -157,11 +180,14 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
 
                 yield* Effect.logError(err)
 
-                ws.send(encodeOutgoingMessage(err))
+                yield* respond(err)
+                yield* this.pushSemaphore.release(1)
                 return
               }
 
-              ws.send(encodeOutgoingMessage(WSMessage.PushAck.make({ requestId })))
+              yield* respond(WSMessage.PushAck.make({ requestId }))
+
+              yield* this.pushSemaphore.release(1)
 
               const createdAt = new Date().toISOString()
 
@@ -172,22 +198,26 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
               const connectedClients = this.ctx.getWebSockets()
 
               // console.debug(`Broadcasting push batch to ${this.subscribedWebSockets.size} clients`)
-
               if (connectedClients.length > 0) {
-                const pullRes = encodeOutgoingMessage(
-                  // TODO refactor to batch api
-                  WSMessage.PullRes.make({
-                    batch: decodedMessage.batch.map((mutationEventEncoded) => ({
-                      mutationEventEncoded,
-                      metadata: Option.some({ createdAt }),
-                    })),
-                    remaining: 0,
-                  }),
-                )
+                // TODO refactor to batch api
+                const pullRes = WSMessage.PullRes.make({
+                  batch: decodedMessage.batch.map((mutationEventEncoded) => ({
+                    mutationEventEncoded,
+                    metadata: Option.some({ createdAt }),
+                  })),
+                  remaining: 0,
+                  requestId: { context: 'push', requestId },
+                })
+                const pullResEnc = encodeOutgoingMessage(pullRes)
+
+                // Only calling once for now.
+                if (options?.onPullRes) {
+                  yield* Effect.tryAll(() => options.onPullRes!(pullRes))
+                }
 
                 // NOTE we're also sending the pullRes to the pushing ws client as a confirmation
                 for (const conn of connectedClients) {
-                  conn.send(pullRes)
+                  conn.send(pullResEnc)
                 }
               }
 
@@ -230,12 +260,15 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
           ws.send(encodeOutgoingMessage(WSMessage.Error.make({ message: error.message, requestId })))
         }
       }).pipe(
-        Effect.withSpan('@livestore/sync-cf:durable-object:webSocketMessage'),
+        Effect.withSpan(`@livestore/sync-cf:durable-object:webSocketMessage:${decodedMessage._tag}`, {
+          attributes: { requestId },
+        }),
         Effect.tapCauseLogPretty,
         Logger.withMinimumLogLevel(LogLevel.Debug),
         Effect.provide(Logger.pretty),
         Effect.runPromise,
       )
+    }
 
     webSocketClose = async (ws: WebSocketClient, code: number, _reason: string, _wasClean: boolean) => {
       // If the client closes the connection, the runtime will invoke the webSocketClose() handler.
@@ -267,7 +300,10 @@ const makeStorage = (ctx: DurableObjectState, env: Env, storeId: string): SyncSt
     Effect.tryPromise({
       try: () => cb(env.DB),
       catch: (error) => new UnexpectedError({ cause: error, payload: { dbName } }),
-    }).pipe(Effect.map((_) => _.results))
+    }).pipe(
+      Effect.map((_) => _.results),
+      Effect.withSpan('@livestore/sync-cf:durable-object:execDb'),
+    )
 
   const getHead: Effect.Effect<EventId.GlobalEventId, UnexpectedError> = Effect.gen(function* () {
     const result = yield* execDb<{ id: EventId.GlobalEventId }>((db) =>
@@ -295,7 +331,7 @@ const makeStorage = (ctx: DurableObjectState, env: Env, storeId: string): SyncSt
         }),
       )
       return events
-    })
+    }).pipe(UnexpectedError.mapToUnexpectedError)
 
   const appendEvents: SyncStorage['appendEvents'] = (batch, createdAt) =>
     Effect.gen(function* () {
@@ -330,7 +366,7 @@ const makeStorage = (ctx: DurableObjectState, env: Env, storeId: string): SyncSt
             .run(),
         )
       }
-    })
+    }).pipe(UnexpectedError.mapToUnexpectedError)
 
   const resetStore = Effect.gen(function* () {
     yield* Effect.promise(() => ctx.storage.deleteAll())

package/src/common/ws-message-types.ts

@@ -23,6 +23,7 @@ export const PullRes = Schema.TaggedStruct('WSMessage.PullRes', {
       metadata: Schema.Option(SyncMetadata),
     }),
   ),
+  requestId: Schema.Struct({ context: Schema.Literal('pull', 'push'), requestId: Schema.String }),
   remaining: Schema.Number,
 })
 
@@ -46,6 +47,8 @@ export const Error = Schema.TaggedStruct('WSMessage.Error', {
   message: Schema.String,
 })
 
+export type Error = typeof Error.Type
+
 export const Ping = Schema.TaggedStruct('WSMessage.Ping', {
   requestId: Schema.Literal('ping'),
 })

package/src/sync-impl/ws-impl.ts

@@ -2,7 +2,8 @@
 
 import type { SyncBackend } from '@livestore/common'
 import { InvalidPullError, InvalidPushError } from '@livestore/common'
-import { LS_DEV } from '@livestore/utils'
+import { EventId } from '@livestore/common/schema'
+import { LS_DEV, shouldNeverHappen } from '@livestore/utils'
 import type { Scope } from '@livestore/utils/effect'
 import {
   Deferred,
@@ -33,10 +34,29 @@ export const makeWsSync = (options: WsSyncOptions): Effect.Effect<SyncBackend<Sy
 
     const { isConnected, incomingMessages, send } = yield* connect(wsUrl)
 
+    /**
+     * We need to account for the scenario where push-caused PullRes message arrive before the pull-caused PullRes message.
+     * i.e. a scenario where the WS connection is created but before the server processed the initial pull, a push from
+     * another client triggers a PullRes message sent to this client which we need to stash until our pull-caused
+     * PullRes message arrives at which point we can combine the stashed events with the pull-caused events and continue.
+     */
+    const stashedPullBatch: WSMessage.PullRes['batch'][number][] = []
+
+    // We currently only support one pull stream for a sync backend.
+    let pullStarted = false
+
     const api = {
       isConnected,
       pull: (args) =>
         Effect.gen(function* () {
+          if (pullStarted) {
+            return shouldNeverHappen(`Pull already started for this sync backend.`)
+          }
+
+          pullStarted = true
+
+          let pullResponseReceived = false
+
           const requestId = nanoid()
           const cursor = Option.getOrUndefined(args)?.cursor.global
 
@@ -48,6 +68,41 @@ export const makeWsSync = (options: WsSyncOptions): Effect.Effect<SyncBackend<Sy
                 ? new InvalidPullError({ message: _.message })
                 : Effect.void,
             ),
+            Stream.filterMap((msg) => {
+              if (msg._tag === 'WSMessage.PullRes') {
+                if (msg.requestId.context === 'pull') {
+                  if (msg.requestId.requestId === requestId) {
+                    pullResponseReceived = true
+
+                    if (stashedPullBatch.length > 0 && msg.remaining === 0) {
+                      const pullResHead = msg.batch.at(-1)?.mutationEventEncoded.id ?? EventId.ROOT.global
+                      // Index where stashed events are greater than pullResHead
+                      const newPartialBatchIndex = stashedPullBatch.findIndex(
+                        (batchItem) => batchItem.mutationEventEncoded.id > pullResHead,
+                      )
+                      const batchWithNewStashedEvents =
+                        newPartialBatchIndex === -1 ? [] : stashedPullBatch.slice(newPartialBatchIndex)
+                      const combinedBatch = [...msg.batch, ...batchWithNewStashedEvents]
+                      return Option.some({ ...msg, batch: combinedBatch, remaining: 0 })
+                    } else {
+                      return Option.some(msg)
+                    }
+                  } else {
+                    // Ignore
+                    return Option.none()
+                  }
+                } else {
+                  if (pullResponseReceived) {
+                    return Option.some(msg)
+                  } else {
+                    stashedPullBatch.push(...msg.batch)
+                    return Option.none()
+                  }
+                }
+              }
+
+              return Option.none()
+            }),
             // This call is mostly here to for type narrowing
             Stream.filter(Schema.is(WSMessage.PullRes)),
           )