@livestore/sync-cf 0.3.0-dev.9 → 0.3.0

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

@@ -1,5 +1,5 @@
-import { makeColumnSpec } from '@livestore/common'
-import { DbSchema, EventId, type MutationEvent } from '@livestore/common/schema'
+import { makeColumnSpec, UnexpectedError } from '@livestore/common'
+import { EventSequenceNumber, type LiveStoreEvent, State } from '@livestore/common/schema'
 import { shouldNeverHappen } from '@livestore/utils'
 import { Effect, Logger, LogLevel, Option, Schema } from '@livestore/utils/effect'
 import { DurableObject } from 'cloudflare:workers'
@@ -8,7 +8,6 @@ import { WSMessage } from '../common/mod.js'
 import type { SyncMetadata } from '../common/ws-message-types.js'
 
 export interface Env {
-  WEBSOCKET_SERVER: DurableObjectNamespace
   DB: D1Database
   ADMIN_SECRET: string
 }
@@ -19,26 +18,41 @@ const encodeOutgoingMessage = Schema.encodeSync(Schema.parseJson(WSMessage.Backe
 const encodeIncomingMessage = Schema.encodeSync(Schema.parseJson(WSMessage.ClientToBackendMessage))
 const decodeIncomingMessage = Schema.decodeUnknownEither(Schema.parseJson(WSMessage.ClientToBackendMessage))
 
-// NOTE actual table name is determined at runtime by `WebSocketServer.dbName`
-export const mutationLogTable = DbSchema.table('__unused', {
-  id: DbSchema.integer({ primaryKey: true, schema: EventId.GlobalEventId }),
-  parentId: DbSchema.integer({ schema: EventId.GlobalEventId }),
-  mutation: DbSchema.text({}),
-  args: DbSchema.text({ schema: Schema.parseJson(Schema.Any) }),
-  /** ISO date format. Currently only used for debugging purposes. */
-  createdAt: DbSchema.text({}),
+export const eventlogTable = State.SQLite.table({
+  // NOTE actual table name is determined at runtime
+  name: 'eventlog_${PERSISTENCE_FORMAT_VERSION}_${storeId}',
+  columns: {
+    seqNum: State.SQLite.integer({ primaryKey: true, schema: EventSequenceNumber.GlobalEventSequenceNumber }),
+    parentSeqNum: State.SQLite.integer({ schema: EventSequenceNumber.GlobalEventSequenceNumber }),
+    name: State.SQLite.text({}),
+    args: State.SQLite.text({ schema: Schema.parseJson(Schema.Any), nullable: true }),
+    /** ISO date format. Currently only used for debugging purposes. */
+    createdAt: State.SQLite.text({}),
+    clientId: State.SQLite.text({}),
+    sessionId: State.SQLite.text({}),
+  },
 })
 
+const WebSocketAttachmentSchema = Schema.parseJson(
+  Schema.Struct({
+    storeId: Schema.String,
+  }),
+)
+
+export const PULL_CHUNK_SIZE = 100
+
 /**
- * Needs to be bumped when the storage format changes (e.g. mutationLogTable schema changes)
+ * Needs to be bumped when the storage format changes (e.g. eventlogTable schema changes)
  *
  * Changing this version number will lead to a "soft reset".
  */
-export const PERSISTENCE_FORMAT_VERSION = 2
+export const PERSISTENCE_FORMAT_VERSION = 7
 
 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) => {
@@ -47,22 +61,21 @@ export type MakeDurableObjectClass = (options?: MakeDurableObjectClassOptions) =
 
 export const makeDurableObject: MakeDurableObjectClass = (options) => {
   return class WebSocketServerBase extends DurableObject<Env> {
-    storage: SyncStorage | undefined = undefined
+    /** Needed to prevent concurrent pushes */
+    private pushSemaphore = Effect.makeSemaphore(1).pipe(Effect.runSync)
 
-    constructor(ctx: DurableObjectState, env: Env) {
-      super(ctx, env)
-    }
+    private currentHead: EventSequenceNumber.GlobalEventSequenceNumber | 'uninitialized' = 'uninitialized'
 
     fetch = async (request: Request) =>
-      Effect.gen(this, function* () {
-        if (this.storage === undefined) {
-          const storeId = getStoreId(request)
-          const dbName = `mutation_log_${PERSISTENCE_FORMAT_VERSION}_${toValidTableName(storeId)}`
-          this.storage = makeStorage(this.ctx, this.env, dbName)
-        }
+      Effect.sync(() => {
+        const storeId = getStoreId(request)
+        const storage = makeStorage(this.ctx, this.env, storeId)
 
         const { 0: client, 1: server } = new WebSocketPair()
 
+        // Since we're using websocket hibernation, we need to remember the storeId for subsequent `webSocketMessage` calls
+        server.serializeAttachment(Schema.encodeSync(WebSocketAttachmentSchema)({ storeId }))
+
         // See https://developers.cloudflare.com/durable-objects/examples/websocket-hibernation-server
 
         this.ctx.acceptWebSocket(server)
@@ -74,8 +87,8 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
           ),
         )
 
-        const colSpec = makeColumnSpec(mutationLogTable.sqliteDef.ast)
-        this.env.DB.exec(`CREATE TABLE IF NOT EXISTS ${this.storage.dbName} (${colSpec}) strict`)
+        const colSpec = makeColumnSpec(eventlogTable.sqliteDef.ast)
+        this.env.DB.exec(`CREATE TABLE IF NOT EXISTS ${storage.dbName} (${colSpec}) strict`)
 
         return new Response(null, {
           status: 101,
@@ -83,114 +96,156 @@ 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) => {
+      console.log('webSocketMessage', message)
+      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
+      const decodedMessage = decodedMessageRes.right
+      const requestId = decodedMessage.requestId
 
-        const storage = this.storage
-
-        if (storage === undefined) {
-          throw new Error('storage not initialized')
-        }
+      return Effect.gen(this, function* () {
+        const { storeId } = yield* Schema.decode(WebSocketAttachmentSchema)(ws.deserializeAttachment())
+        const storage = makeStorage(this.ctx, this.env, storeId)
 
         try {
           switch (decodedMessage._tag) {
+            // TODO allow pulling concurrently to not block incoming push requests
             case 'WSMessage.PullReq': {
               if (options?.onPull) {
                 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* Effect.promise(() => storage.getEvents(cursor)))]
-
-              // NOTE we want to make sure the WS server responds at least once with `InitRes` even if `events` is empty
-              while (true) {
-                const events = remainingEvents.splice(0, CHUNK_SIZE)
-
-                ws.send(
-                  encodeOutgoingMessage(
-                    WSMessage.PullRes.make({ events, remaining: remainingEvents.length, requestId }),
-                  ),
-                )
-
-                if (remainingEvents.length === 0) {
-                  break
-                }
+              const remainingEvents = yield* storage.getEvents(cursor)
+
+              // Send at least one response, even if there are no events
+              const batches =
+                remainingEvents.length === 0
+                  ? [[]]
+                  : 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) * PULL_CHUNK_SIZE)
+                yield* respond(WSMessage.PullRes.make({ batch, remaining, requestId: { context: 'pull', requestId } }))
               }
 
               break
             }
             case 'WSMessage.PushReq': {
+              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) {
+                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 latestEvent = yield* Effect.promise(() => storage.getLatestEvent())
-              const expectedParentId = latestEvent?.id ?? EventId.ROOT.global
+              // const expectedParentNum = yield* storage.getHead
+
+              let currentHead: EventSequenceNumber.GlobalEventSequenceNumber
+              if (this.currentHead === 'uninitialized') {
+                const currentHeadFromStorage = yield* Effect.promise(() => this.ctx.storage.get('currentHead'))
+                // console.log('currentHeadFromStorage', currentHeadFromStorage)
+                if (currentHeadFromStorage === undefined) {
+                  // console.log('currentHeadFromStorage is null, getting from D1')
+                  // currentHead = yield* storage.getHead
+                  // console.log('currentHeadFromStorage is null, using root')
+                  currentHead = EventSequenceNumber.ROOT.global
+                } else {
+                  currentHead = currentHeadFromStorage as EventSequenceNumber.GlobalEventSequenceNumber
+                }
+              } else {
+                // console.log('currentHead is already initialized', this.currentHead)
+                currentHead = this.currentHead
+              }
 
-              let i = 0
-              for (const mutationEventEncoded of decodedMessage.batch) {
-                if (mutationEventEncoded.parentId !== expectedParentId + i) {
-                  const err = WSMessage.Error.make({
-                    message: `Invalid parent id. Received ${mutationEventEncoded.parentId} but expected ${expectedParentId}`,
-                    requestId,
-                  })
+              // TODO handle clientId unique conflict
+              // Validate the batch
+              const firstEvent = decodedMessage.batch[0]!
+              if (firstEvent.parentSeqNum !== currentHead) {
+                const err = WSMessage.Error.make({
+                  message: `Invalid parent event number. Received e${firstEvent.parentSeqNum} but expected e${currentHead}`,
+                  requestId,
+                })
 
-                  yield* Effect.fail(err).pipe(Effect.ignoreLogged)
+                yield* Effect.logError(err)
 
-                  ws.send(encodeOutgoingMessage(err))
-                  return
-                }
+                yield* respond(err)
+                yield* this.pushSemaphore.release(1)
+                return
+              }
 
-                // TODO handle clientId unique conflict
+              yield* respond(WSMessage.PushAck.make({ requestId }))
 
-                const createdAt = new Date().toISOString()
+              const createdAt = new Date().toISOString()
 
-                // NOTE we're currently not blocking on this to allow broadcasting right away
-                const storePromise = storage.appendEvent(mutationEventEncoded, createdAt)
+              // NOTE we're not waiting for this to complete yet to allow the broadcast to happen right away
+              // while letting the async storage write happen in the background
+              const storeFiber = yield* storage.appendEvents(decodedMessage.batch, createdAt).pipe(Effect.fork)
 
-                ws.send(
-                  encodeOutgoingMessage(WSMessage.PushAck.make({ mutationId: mutationEventEncoded.id, requestId })),
-                )
+              this.currentHead = decodedMessage.batch.at(-1)!.seqNum
+              yield* Effect.promise(() => this.ctx.storage.put('currentHead', this.currentHead))
 
-                // console.debug(`Broadcasting mutation event to ${this.subscribedWebSockets.size} clients`)
+              yield* this.pushSemaphore.release(1)
 
-                const connectedClients = this.ctx.getWebSockets()
+              const connectedClients = this.ctx.getWebSockets()
 
-                if (connectedClients.length > 0) {
-                  const broadcastMessage = encodeOutgoingMessage(
-                    // TODO refactor to batch api
-                    WSMessage.PushBroadcast.make({
-                      mutationEventEncoded,
-                      metadata: Option.some({ createdAt }),
-                    }),
-                  )
+              // console.debug(`Broadcasting push batch to ${this.subscribedWebSockets.size} clients`)
+              if (connectedClients.length > 0) {
+                // TODO refactor to batch api
+                const pullRes = WSMessage.PullRes.make({
+                  batch: decodedMessage.batch.map((eventEncoded) => ({
+                    eventEncoded,
+                    metadata: Option.some({ createdAt }),
+                  })),
+                  remaining: 0,
+                  requestId: { context: 'push', requestId },
+                })
+                const pullResEnc = encodeOutgoingMessage(pullRes)
 
-                  for (const conn of connectedClients) {
-                    console.log('Broadcasting to client', conn === ws ? 'self' : 'other')
-                    // if (conn !== ws) {
-                    conn.send(broadcastMessage)
-                    // }
-                  }
+                // Only calling once for now.
+                if (options?.onPullRes) {
+                  yield* Effect.tryAll(() => options.onPullRes!(pullRes))
                 }
 
-                yield* Effect.promise(() => storePromise)
-
-                i++
+                // NOTE we're also sending the pullRes to the pushing ws client as a confirmation
+                for (const conn of connectedClients) {
+                  conn.send(pullResEnc)
+                }
               }
 
+              // Wait for the storage write to complete before finishing this request
+              yield* storeFiber
+
               break
             }
             case 'WSMessage.AdminResetRoomReq': {
@@ -199,7 +254,7 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
                 return
               }
 
-              yield* Effect.promise(() => storage.resetRoom())
+              yield* storage.resetStore
               ws.send(encodeOutgoingMessage(WSMessage.AdminResetRoomRes.make({ requestId })))
 
               break
@@ -227,12 +282,20 @@ 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,
+        Effect.tapErrorCause((cause) =>
+          Effect.sync(() =>
+            ws.send(encodeOutgoingMessage(WSMessage.Error.make({ message: cause.toString(), requestId }))),
+          ),
+        ),
         Logger.withMinimumLogLevel(LogLevel.Debug),
-        Effect.provide(Logger.pretty),
+        Effect.provide(Logger.prettyWithThread('durable-object')),
         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.
@@ -243,60 +306,109 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
 
 type SyncStorage = {
   dbName: string
-  getLatestEvent: () => Promise<MutationEvent.AnyEncodedGlobal | undefined>
+  // getHead: Effect.Effect<EventSequenceNumber.GlobalEventSequenceNumber, UnexpectedError>
   getEvents: (
     cursor: number | undefined,
-  ) => Promise<
-    ReadonlyArray<{ mutationEventEncoded: MutationEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> }>
+  ) => Effect.Effect<
+    ReadonlyArray<{ eventEncoded: LiveStoreEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> }>,
+    UnexpectedError
   >
-  appendEvent: (event: MutationEvent.AnyEncodedGlobal, createdAt: string) => Promise<void>
-  resetRoom: () => Promise<void>
+  appendEvents: (
+    batch: ReadonlyArray<LiveStoreEvent.AnyEncodedGlobal>,
+    createdAt: string,
+  ) => Effect.Effect<void, UnexpectedError>
+  resetStore: Effect.Effect<void, UnexpectedError>
 }
 
-const makeStorage = (ctx: DurableObjectState, env: Env, dbName: string): SyncStorage => {
-  const getLatestEvent = async (): Promise<MutationEvent.AnyEncodedGlobal | undefined> => {
-    const rawEvents = await env.DB.prepare(`SELECT * FROM ${dbName} ORDER BY id DESC LIMIT 1`).all()
-    if (rawEvents.error) {
-      throw new Error(rawEvents.error)
-    }
-    const events = Schema.decodeUnknownSync(Schema.Array(mutationLogTable.schema))(rawEvents.results)
+const makeStorage = (ctx: DurableObjectState, env: Env, storeId: string): SyncStorage => {
+  const dbName = `eventlog_${PERSISTENCE_FORMAT_VERSION}_${toValidTableName(storeId)}`
 
-    return events[0]
-  }
-
-  const getEvents = async (
-    cursor: number | undefined,
-  ): Promise<
-    ReadonlyArray<{ mutationEventEncoded: MutationEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> }>
-  > => {
-    const whereClause = cursor === undefined ? '' : `WHERE id > ${cursor}`
-    const sql = `SELECT * FROM ${dbName} ${whereClause} ORDER BY id ASC`
-    // TODO handle case where `cursor` was not found
-    const rawEvents = await env.DB.prepare(sql).all()
-    if (rawEvents.error) {
-      throw new Error(rawEvents.error)
-    }
-    const events = Schema.decodeUnknownSync(Schema.Array(mutationLogTable.schema))(rawEvents.results).map(
-      ({ createdAt, ...mutationEventEncoded }) => ({
-        mutationEventEncoded,
-        metadata: Option.some({ createdAt }),
-      }),
+  const execDb = <T>(cb: (db: D1Database) => Promise<D1Result<T>>) =>
+    Effect.tryPromise({
+      try: () => cb(env.DB),
+      catch: (error) => new UnexpectedError({ cause: error, payload: { dbName } }),
+    }).pipe(
+      Effect.map((_) => _.results),
+      Effect.withSpan('@livestore/sync-cf:durable-object:execDb'),
     )
-    return events
-  }
 
-  const appendEvent = async (event: MutationEvent.AnyEncodedGlobal, createdAt: string) => {
-    const sql = `INSERT INTO ${dbName} (id, parentId, args, mutation, createdAt) VALUES (?, ?, ?, ?, ?)`
-    await env.DB.prepare(sql)
-      .bind(event.id, event.parentId, JSON.stringify(event.args), event.mutation, createdAt)
-      .run()
-  }
+  // const getHead: Effect.Effect<EventSequenceNumber.GlobalEventSequenceNumber, UnexpectedError> = Effect.gen(
+  //   function* () {
+  //     const result = yield* execDb<{ seqNum: EventSequenceNumber.GlobalEventSequenceNumber }>((db) =>
+  //       db.prepare(`SELECT seqNum FROM ${dbName} ORDER BY seqNum DESC LIMIT 1`).all(),
+  //     )
 
-  const resetRoom = async () => {
-    await ctx.storage.deleteAll()
-  }
+  //     return result[0]?.seqNum ?? EventSequenceNumber.ROOT.global
+  //   },
+  // ).pipe(UnexpectedError.mapToUnexpectedError)
 
-  return { dbName, getLatestEvent, getEvents, appendEvent, resetRoom }
+  const getEvents = (
+    cursor: number | undefined,
+  ): Effect.Effect<
+    ReadonlyArray<{ eventEncoded: LiveStoreEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> }>,
+    UnexpectedError
+  > =>
+    Effect.gen(function* () {
+      const whereClause = cursor === undefined ? '' : `WHERE seqNum > ${cursor}`
+      const sql = `SELECT * FROM ${dbName} ${whereClause} ORDER BY seqNum ASC`
+      // TODO handle case where `cursor` was not found
+      const rawEvents = yield* execDb((db) => db.prepare(sql).all())
+      const events = Schema.decodeUnknownSync(Schema.Array(eventlogTable.rowSchema))(rawEvents).map(
+        ({ createdAt, ...eventEncoded }) => ({
+          eventEncoded,
+          metadata: Option.some({ createdAt }),
+        }),
+      )
+      return events
+    }).pipe(UnexpectedError.mapToUnexpectedError)
+
+  const appendEvents: SyncStorage['appendEvents'] = (batch, createdAt) =>
+    Effect.gen(function* () {
+      // If there are no events, do nothing.
+      if (batch.length === 0) return
+
+      // CF D1 limits:
+      // Maximum bound parameters per query	100, Maximum arguments per SQL function	32
+      // Thus we need to split the batch into chunks of max (100/7=)14 events each.
+      const CHUNK_SIZE = 14
+
+      for (let i = 0; i < batch.length; i += CHUNK_SIZE) {
+        const chunk = batch.slice(i, i + CHUNK_SIZE)
+
+        // Create a list of placeholders ("(?, ?, ?, ?, ?, ?, ?)"), corresponding to each event.
+        const valuesPlaceholders = chunk.map(() => '(?, ?, ?, ?, ?, ?, ?)').join(', ')
+        const sql = `INSERT INTO ${dbName} (seqNum, parentSeqNum, args, name, createdAt, clientId, sessionId) VALUES ${valuesPlaceholders}`
+        // Flatten the event properties into a parameters array.
+        const params = chunk.flatMap((event) => [
+          event.seqNum,
+          event.parentSeqNum,
+          event.args === undefined ? null : JSON.stringify(event.args),
+          event.name,
+          createdAt,
+          event.clientId,
+          event.sessionId,
+        ])
+
+        yield* execDb((db) =>
+          db
+            .prepare(sql)
+            .bind(...params)
+            .run(),
+        )
+      }
+    }).pipe(UnexpectedError.mapToUnexpectedError)
+
+  const resetStore = Effect.gen(function* () {
+    yield* Effect.promise(() => ctx.storage.deleteAll())
+  }).pipe(UnexpectedError.mapToUnexpectedError)
+
+  return {
+    dbName,
+    // getHead,
+    getEvents,
+    appendEvents,
+    resetStore,
+  }
 }
 
 const getStoreId = (request: Request) => {