fix(collab): preserve initial sync delivery under backpressure

server/collab/socket.ts

@@ -55,6 +55,7 @@ import type { DbClient } from '../db/client'
 import { jsonResponse } from '../http'
 import { BranchGoneError } from './relayBranches'
 import type { CollabRelay, RelayDoc } from './relay'
+import { createSocketFlow } from './socketFlow'
 
 export { SITE_SOCKET_PATH }
 
@@ -244,6 +245,12 @@ interface CollabPublisher {
  * awareness instance.
  */
 export function createCollabSocketLayer(relay: CollabRelay) {
+  const flows = new WeakMap<ServerWebSocket<CollabSocketData>, ReturnType<typeof createSocketFlow>>()
+  function flowFor(ws: ServerWebSocket<CollabSocketData>) {
+    let flow = flows.get(ws)
+    if (!flow) { flow = createSocketFlow(ws); flows.set(ws, flow) }
+    return flow
+  }
   let publisher: CollabPublisher | null = null
   const presenceDoc = new Y.Doc()
   const awareness = new awarenessProtocol.Awareness(presenceDoc)
@@ -268,12 +275,12 @@ export function createCollabSocketLayer(relay: CollabRelay) {
   })
 
   /** Tell one connection its doc was dropped, and why. */
-  function sendReset(
+  async function sendReset(
     ws: ServerWebSocket<CollabSocketData>,
     docId: string,
     reason: ResetReason,
-  ): void {
-    ws.send(encodeCollabFrame(docId, '', FRAME_RESET, encodeResetPayload(reason)))
+  ): Promise<void> {
+    await flowFor(ws).send(encodeCollabFrame(docId, '', FRAME_RESET, encodeResetPayload(reason)))
   }
 
   /**
@@ -289,7 +296,7 @@ export function createCollabSocketLayer(relay: CollabRelay) {
       // parseCollabDocId or the relay. Deliberately ungated — a read-only
       // viewer needs to know its socket is alive exactly as much as a writer.
       if (frame.frameType === FRAME_PING) {
-        ws.send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_PONG, new Uint8Array()))
+        await flowFor(ws).send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_PONG, new Uint8Array()))
         return
       }
 
@@ -318,7 +325,7 @@ export function createCollabSocketLayer(relay: CollabRelay) {
         )
         // publish() excludes nobody server-side; the sender's own state is
         // already local — awareness re-application is idempotent.
-        ws.send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_AWARENESS, frame.payload))
+        await flowFor(ws).send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_AWARENESS, frame.payload))
         return
       }
 
@@ -332,7 +339,7 @@ export function createCollabSocketLayer(relay: CollabRelay) {
         // structs the server will never receive, every later update queues
         // behind them as pending, and the screen shows edits that can never
         // publish. A visible revert is strictly better.
-        sendReset(ws, frame.docId, 'oversize')
+        await sendReset(ws, frame.docId, 'oversize')
         return
       }
 
@@ -372,14 +379,14 @@ export function createCollabSocketLayer(relay: CollabRelay) {
         // before any inbound frame has taught it a generation.
         const serverIsEmpty = Y.encodeStateVector(doc).byteLength === 1
         if (messageType !== SYNC_STEP_1 && !serverIsEmpty) {
-          sendReset(ws, frame.docId, 'stale')
+          await sendReset(ws, frame.docId, 'stale')
           return
         }
       } else if (frame.generation !== generation) {
         console.warn(
           `[collab] stale generation for ${frame.docId} from ${ws.data.userId}`,
         )
-        sendReset(ws, frame.docId, 'stale')
+        await sendReset(ws, frame.docId, 'stale')
         return
       }
 
@@ -402,7 +409,7 @@ export function createCollabSocketLayer(relay: CollabRelay) {
           // The sender's local doc holds the forbidden change — a TARGETED
           // reset makes their client rebind and reseed from the server,
           // reverting it everywhere (including their own screen).
-          sendReset(ws, frame.docId, 'refused')
+          await sendReset(ws, frame.docId, 'refused')
           return
         }
         Y.applyUpdate(doc, update, ws)
@@ -428,64 +435,73 @@ export function createCollabSocketLayer(relay: CollabRelay) {
         ws.data.probedDocs.add(frame.docId)
         const probe = encoding.createEncoder()
         syncProtocol.writeSyncStep1(probe, doc)
-        ws.send(encodeCollabFrame(frame.docId, generation, FRAME_SYNC, encoding.toUint8Array(probe)))
+        await flowFor(ws).send(encodeCollabFrame(frame.docId, generation, FRAME_SYNC, encoding.toUint8Array(probe)))
       }
 
       const decoder = decoding.createDecoder(frame.payload)
       const encoder = encoding.createEncoder()
       syncProtocol.readSyncMessage(decoder, encoder, doc, ws)
       if (encoding.length(encoder) > 0) {
-        ws.send(encodeCollabFrame(frame.docId, generation, FRAME_SYNC, encoding.toUint8Array(encoder)))
+        await flowFor(ws).send(encodeCollabFrame(frame.docId, generation, FRAME_SYNC, encoding.toUint8Array(encoder)))
       }
   }
 
   const handlers: WebSocketHandler<CollabSocketData> = {
     // Transport-level ceiling — the per-frame-type caps in `message` are the
     // fine-grained guards; this stops oversized frames before they buffer.
     maxPayloadLength: MAX_SYNC_PAYLOAD_BYTES + 1024,
+    // Pub/sub frames also need recovery rather than silent delivery loss.
+    closeOnBackpressureLimit: true,
 
-    open(ws: ServerWebSocket<CollabSocketData>) {
-      ws.subscribe(docTopic(PRESENCE_DOC_ID))
-      // Late joiners need the current presence roster.
-      const known = [...awareness.getStates().keys()]
-      if (known.length > 0) {
-        const update = awarenessProtocol.encodeAwarenessUpdate(awareness, known)
-        ws.send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_AWARENESS, update))
-      }
+    async open(ws: ServerWebSocket<CollabSocketData>) {
+      await flowFor(ws).run(async () => {
+        ws.subscribe(docTopic(PRESENCE_DOC_ID))
+        // Late joiners need the current presence roster.
+        const known = [...awareness.getStates().keys()]
+        if (known.length > 0) {
+          const update = awarenessProtocol.encodeAwarenessUpdate(awareness, known)
+          await flowFor(ws).send(encodeCollabFrame(PRESENCE_DOC_ID, '', FRAME_AWARENESS, update))
+        }
+      })
     },
 
+    drain(ws) { flowFor(ws).drain() },
+
     async message(ws: ServerWebSocket<CollabSocketData>, raw: string | Buffer) {
       if (typeof raw === 'string') return // binary protocol only
-      let frame: CollabFrame | null = null
-      try {
-        frame = decodeCollabFrame(new Uint8Array(raw))
-        await dispatchFrame(ws, frame)
-      } catch (err) {
-        if (err instanceof BranchGoneError && frame) {
-          // The doc's branch was deleted. The client must leave the branch,
-          // not rebind — rebinding would loop through this refusal forever.
-          try {
-            sendReset(ws, frame.docId, 'gone')
-          } catch (_sendErr) {
-            // Socket already closing — nothing to recover.
+      await flowFor(ws).run(async () => {
+        let frame: CollabFrame | null = null
+        try {
+          frame = decodeCollabFrame(new Uint8Array(raw))
+          await dispatchFrame(ws, frame)
+        } catch (err) {
+          if (err instanceof BranchGoneError && frame) {
+            // The doc's branch was deleted. The client must leave the branch,
+            // not rebind — rebinding would loop through this refusal forever.
+            try {
+              await sendReset(ws, frame.docId, 'gone')
+            } catch (_sendErr) {
+              // Socket already closing — nothing to recover.
+            }
+            return
           }
-          return
-        }
-        console.error('[collab] socket message handler failed:', err)
-        // A sync-write frame whose guard/apply threw left the sender's local
-        // doc diverged from the authoritative one — reset it so their client
-        // rebinds and reseeds. Awareness/malformed frames just get dropped.
-        if (frame && frame.frameType === FRAME_SYNC && parseCollabDocId(frame.docId)) {
-          try {
-            sendReset(ws, frame.docId, 'refused')
-          } catch (_sendErr) {
-            // Socket already closing — nothing to recover.
+          console.error('[collab] socket message handler failed:', err)
+          // A sync-write frame whose guard/apply threw left the sender's local
+          // doc diverged from the authoritative one — reset it so their client
+          // rebinds and reseeds. Awareness/malformed frames just get dropped.
+          if (frame && frame.frameType === FRAME_SYNC && parseCollabDocId(frame.docId)) {
+            try {
+              await sendReset(ws, frame.docId, 'refused')
+            } catch (_sendErr) {
+              // Socket already closing — nothing to recover.
+            }
           }
         }
-      }
+      })
     },
 
     close(ws: ServerWebSocket<CollabSocketData>) {
+      flowFor(ws).close()
       for (const docId of ws.data.boundDocs) relay.release(docId)
       ws.data.boundDocs.clear()
       if (ws.data.awarenessClients.size > 0) {

server/collab/socketFlow.ts

@@ -0,0 +1,46 @@
+/** Serialize incoming work and stop producing replies while Bun buffers a send. */
+interface FlowSocket {
+  send(frame: Uint8Array): number
+  close(code: number, reason: string): void
+}
+
+export function createSocketFlow(socket: FlowSocket) {
+  let closed = false
+  let tail = Promise.resolve()
+  let resume: (() => void) | null = null
+
+  function drain(): void {
+    const pending = resume
+    resume = null
+    pending?.()
+  }
+
+  return {
+    run(task: () => Promise<void>): Promise<void> {
+      const next = tail.then(async () => {
+        if (!closed) await task()
+      })
+      // A rejected request must not poison subsequent socket requests.
+      tail = next.catch((_err) => {})
+      return next
+    },
+    async send(frame: Uint8Array): Promise<void> {
+      if (closed) return
+      const result = socket.send(frame)
+      if (result === 0) {
+        // The frame was dropped: force the provider's state-vector recovery.
+        closed = true
+        drain()
+        socket.close(1011, 'Collaboration delivery failed')
+      } else if (result === -1) {
+        // Bun already queued this frame. Sending it again would duplicate it.
+        await new Promise<void>((resolve) => { resume = resolve })
+      }
+    },
+    drain,
+    close(): void {
+      closed = true
+      drain()
+    },
+  }
+}

src/__tests__/server/collabSocketFlow.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, it } from 'bun:test'
+import { createSocketFlow } from '../../../server/collab/socketFlow'
+
+describe('collab socket backpressure', () => {
+  it('waits for drain before processing the next request without resending the queued frame', async () => {
+    const sent: number[] = []
+    const flow = createSocketFlow({ send: (data) => { sent.push(data[0]!); return sent.length === 1 ? -1 : 1 }, close: () => {} })
+    const first = flow.run(() => flow.send(new Uint8Array([1])))
+    const second = flow.run(() => flow.send(new Uint8Array([2])))
+    await Promise.resolve()
+    expect(sent).toEqual([1])
+    flow.drain()
+    await Promise.all([first, second])
+    expect(sent).toEqual([1, 2])
+  })
+
+  it('settles queued work on disconnect without sending it to the closed socket', async () => {
+    const sent: number[] = []
+    const flow = createSocketFlow({ send: (data) => { sent.push(data[0]!); return -1 }, close: () => {} })
+    const first = flow.run(() => flow.send(new Uint8Array([1])))
+    const second = flow.run(() => flow.send(new Uint8Array([2])))
+    await Promise.resolve()
+    flow.close()
+    await Promise.all([first, second])
+    expect(sent).toEqual([1])
+  })
+
+  it('forces reconnect when Bun reports a dropped frame', async () => {
+    let closed = 0
+    let sends = 0
+    const flow = createSocketFlow({ send: () => { sends++; return 0 }, close: () => { closed++ } })
+    await flow.run(() => flow.send(new Uint8Array([1])))
+    await flow.run(() => flow.send(new Uint8Array([2])))
+    expect(closed).toBe(1)
+    expect(sends).toBe(1)
+  })
+})