Merge pull request #6048 from simstudioai/collab-seed-atomic

fix(realtime): make collab-doc seeding atomic to close split-brain window

Author: waleedlatif1

apps/realtime/src/handlers/file-doc-store.test.ts

@@ -71,9 +71,21 @@ function makeClient(): any {
       b().kv.delete(key)
       return 1
     },
-    // Compare-and-delete Lua (RELEASE_LOCK_SCRIPT): del only if the stored value matches the token.
-    eval: async (_script: string, opts: { keys: string[]; arguments: string[] }) => {
+    eval: async (script: string, opts: { keys: string[]; arguments: string[] }) => {
       const [key] = opts.keys
+      // Atomic seed-if-empty (SEED_IF_EMPTY_SCRIPT): append the entry iff the stream is empty, in one
+      // synchronous step — mirroring Redis's atomic Lua execution, so two concurrent evals can never both
+      // append (the second sees a non-empty stream).
+      if (script.includes('xlen')) {
+        const [field, value] = opts.arguments
+        const arr = b().streams.get(key) ?? []
+        if (arr.length > 0) return 0
+        const id = `${++b().seq}-0`
+        arr.push({ id, message: { [field]: value } })
+        b().streams.set(key, arr)
+        return 1
+      }
+      // Compare-and-delete Lua (RELEASE_LOCK_SCRIPT): del only if the stored value matches the token.
       const [token] = opts.arguments
       if (b().kv.get(key) === token) {
         b().kv.delete(key)
@@ -277,4 +289,106 @@ describe('FileDocStore', () => {
     expect(doc.getText('body').toString()).toBe('')
     doc.destroy()
   })
+
+  it('seedIfEmpty writes the seed once and reports it, then refuses a non-empty stream', async () => {
+    const a = await newStore()
+    expect(await a.seedIfEmpty(NAME, updateFor('first'))).toBe(true)
+    // A second seed attempt (any task) must be refused — the stream already holds content.
+    const b = await newStore()
+    expect(await b.seedIfEmpty(NAME, updateFor('second'))).toBe(false)
+    const doc = new Y.Doc()
+    Y.applyUpdate(doc, (await a.getStreamState(NAME))!)
+    expect(doc.getText('body').toString()).toBe('first')
+    doc.destroy()
+  })
+
+  it('atomic seed prevents split-brain even when the seed lock expired mid-seed', async () => {
+    // The exact split-brain precondition from the concurrency audit: the seed lock is only an efficiency
+    // optimization, so if it lapses (TTL) while the stream is still empty, TWO tasks can both hold a
+    // token and both try to seed with DIFFERENT docs (distinct Yjs client ids). The atomic seedIfEmpty
+    // must still let only one land — otherwise the union duplicates content.
+    const a = await newStore()
+    const b = await newStore()
+    const tokenA = await a.shouldSeed(NAME)
+    expect(tokenA).toBeTruthy()
+    // Simulate A's lock expiring mid-seed so B also wins the freed lock over a still-empty stream.
+    state.backing!.kv.delete(`filedoc:seedlock:${NAME}`)
+    const tokenB = await b.shouldSeed(NAME)
+    expect(tokenB).toBeTruthy()
+    // Both tasks now race to seed with distinct client ids.
+    const [seededA, seededB] = await Promise.all([
+      a.seedIfEmpty(NAME, updateFor('SEED-A')),
+      b.seedIfEmpty(NAME, updateFor('SEED-B')),
+    ])
+    expect([seededA, seededB].filter(Boolean)).toHaveLength(1)
+    // Exactly one seed is in the stream — the reconstructed text is a single seed, never a duplicated
+    // union of both (e.g. 'SEED-ASEED-B').
+    const doc = new Y.Doc()
+    Y.applyUpdate(doc, (await a.getStreamState(NAME))!)
+    expect(['SEED-A', 'SEED-B']).toContain(doc.getText('body').toString())
+    doc.destroy()
+  })
+
+  it('a peer edit published during attachRoom catch-up is not lost', async () => {
+    // Author two INCREMENTAL edits from one doc so they converge to 'basepeer' (not an independent union).
+    const author = new Y.Doc()
+    const updates: Uint8Array[] = []
+    author.on('update', (u: Uint8Array) => updates.push(u))
+    author.getText('body').insert(0, 'base')
+    author.getText('body').insert(4, 'peer')
+
+    const a = await newStore()
+    a.publish(NAME, updates[0]) // 'base'
+    await vi.waitFor(async () => expect(await a.getStreamState(NAME)).not.toBeNull())
+
+    // Task B attaches; while its synchronous catch-up runs, task A publishes the second edit. The tailer
+    // resumes from the id catch-up stopped at, so the edit converges rather than falling into a gap.
+    const b = await newStore()
+    const bDoc = new Y.Doc()
+    const attach = b.attachRoom(NAME, bDoc)
+    a.publish(NAME, updates[1]) // 'peer' appended
+    await attach
+    await vi.waitFor(() => expect(bDoc.getText('body').toString()).toBe('basepeer'), {
+      timeout: 2000,
+    })
+    bDoc.destroy()
+    author.destroy()
+  })
+
+  it('concurrent compaction on two tasks preserves the full document', async () => {
+    const streamKey = `filedoc:stream:${NAME}`
+    const noop = Buffer.from(Y.encodeStateAsUpdate(new Y.Doc())).toString('base64')
+
+    // Two peer edits neither compacting task has integrated (id > each task's lastId).
+    const peerDoc = new Y.Doc()
+    const peerUpdates: Uint8Array[] = []
+    peerDoc.on('update', (u: Uint8Array) => peerUpdates.push(u))
+    peerDoc.getText('body').insert(0, 'PEER1')
+    peerDoc.getText('body').insert(5, 'PEER2')
+
+    const entries = Array.from({ length: 400 }, (_, i) => ({
+      id: `${i + 1}-0`,
+      message: { u: noop },
+    }))
+    entries.push({ id: '401-0', message: { u: Buffer.from(peerUpdates[0]).toString('base64') } })
+    entries.push({ id: '402-0', message: { u: Buffer.from(peerUpdates[1]).toString('base64') } })
+    state.backing!.streams.set(streamKey, entries)
+    state.backing!.seq = 402
+
+    // Two tasks whose local docs lag at DIFFERENT points (400 and 401): both cross the threshold and
+    // compact concurrently. Each must only trim what its own snapshot subsumes, so the union of both
+    // snapshots plus the un-integrated peer entries still reconstructs the whole doc.
+    const a = await newStore()
+    const b = await newStore()
+    const docA = new Y.Doc()
+    Y.applyUpdate(docA, peerUpdates[0]) // A integrated up to 401
+    ;(a as any).rooms.set(NAME, { doc: docA, lastId: '401-0', publishes: 0 })
+    ;(b as any).rooms.set(NAME, { doc: new Y.Doc(), lastId: '400-0', publishes: 0 })
+    await Promise.all([(a as any).maybeCompact(NAME), (b as any).maybeCompact(NAME)])
+
+    const doc = new Y.Doc()
+    Y.applyUpdate(doc, (await a.getStreamState(NAME))!)
+    expect(doc.getText('body').toString()).toBe('PEER1PEER2')
+    doc.destroy()
+  })
 })

apps/realtime/src/handlers/file-doc-store.ts

@@ -22,8 +22,10 @@
  *   opens a file, so a late-joining task (the normal case under autoscaling) loads the current shared
  *   state before its first client syncs. Catch-up + tail are seamless: the tailer resumes from the
  *   exact id catch-up stopped at.
- * - {@link shouldSeed} coordinates the one-time seed with a Redis lock AND an empty-stream check, so
- *   exactly one task ever writes the seed cluster-wide (the fix for split-brain).
+ * - The one-time seed is written via the atomic {@link seedIfEmpty} (append-iff-empty in one Redis
+ *   step), so exactly one task ever writes the seed cluster-wide (the fix for split-brain) — even if two
+ *   tasks race. {@link shouldSeed} is a Redis lock + empty-stream check layered on top ONLY as an
+ *   efficiency gate (so tasks don't all run the seed fetch); correctness does not depend on it.
  *
  * When `REDIS_URL` is unset (single-pod dev) the store is DISABLED and every method degrades to the
  * relay's original single-replica behavior: seed locally, no stream, no tailer.
@@ -48,6 +50,17 @@ const logger = createLogger('FileDocStore')
 const RELEASE_LOCK_SCRIPT =
   "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end"
 
+/**
+ * Atomic seed: append the seed entry ONLY if the stream is still empty, in one Redis-side step. This is
+ * the real split-brain guard — two tasks racing (even both past an expired seed lock) can never both
+ * write a seed (each would mint a distinct Yjs client id → duplicated content), because the emptiness
+ * check and the append happen atomically with no check-then-append window. The seed lock is only an
+ * efficiency optimization (avoid two seed fetches); correctness does not depend on it staying held.
+ * Returns 1 if THIS call wrote the seed, 0 if the stream already had content.
+ */
+const SEED_IF_EMPTY_SCRIPT =
+  "if redis.call('xlen', KEYS[1]) == 0 then redis.call('xadd', KEYS[1], '*', ARGV[1], ARGV[2]); return 1 else return 0 end"
+
 /**
  * The transaction origin the store stamps on updates it applies from the stream. The relay's
  * `doc.on('update')` handler uses it to distinguish an update that ARRIVED from a peer (fan out to
@@ -101,12 +114,12 @@ const COMPACT_LOCK_TTL_MS = 10_000
 /** Retry a failed stream append this many times before giving up, so a transient Redis blip doesn't
  * silently drop an edit from the shared log (which no peer would then ever see). */
 const PUBLISH_MAX_RETRIES = 3
-/** The seed lock spans the app seed fetch (hard-bounded at `seedRequestMs = 8s`) + the apply + the
- * AWAITED seed publish. The margin comfortably exceeds the fetch bound so the lock does not expire
- * mid-seed, while staying at the client readiness deadline (12s) so a dead seeder's lock frees when
- * clients would recover anyway. Double-seed is prevented regardless of the margin: the seeder publishes
- * the seed to the stream BEFORE releasing the lock, so any later seeder's {@link streamHasContent} fence
- * sees it. */
+/** The seed lock spans the app seed fetch (hard-bounded at `seedRequestMs = 8s`) + the atomic seed
+ * append. It is only an EFFICIENCY optimization — it stops two tasks both running the seed fetch — and is
+ * sized to comfortably exceed the fetch bound while staying near the client readiness deadline (12s) so a
+ * dead seeder's lock frees when clients would recover anyway. Double-seed is prevented even if the lock
+ * expires mid-seed, because the seed is written via the atomic {@link SEED_IF_EMPTY_SCRIPT}
+ * (append-iff-empty), NOT the lock — correctness never depends on the lock staying held. */
 const SEED_LOCK_TTL_MS = FILE_DOC_TIMEOUTS.seedRequestMs + 4_000
 /** How long a stream survives with no heartbeat — long enough that an occupied-but-idle doc never
  * loses its shared state (the heartbeat refreshes it while any task holds the room). */
@@ -279,11 +292,43 @@ export class FileDocStore {
   }
 
   /**
-   * Whether the file's stream already holds content — fences a seed apply against a peer that seeded
-   * while this task held a (possibly stale) lock, so we never write a second seed (split-brain). Both
-   * callers treat `true` as "do not seed", so this fails CLOSED: a Redis `xLen` error returns `true`
-   * (cannot confirm the stream is empty → do not risk a double-seed). `false` only when genuinely empty,
-   * or when disabled (single-replica, where seeding locally is always correct).
+   * Atomically seed the stream iff it is still empty (see {@link SEED_IF_EMPTY_SCRIPT}). This is what
+   * actually prevents split-brain double-seeding: the emptiness check and the append happen in one
+   * Redis-side step, so — unlike a separate {@link streamHasContent} fence + {@link publishAndWait} —
+   * there is no check-then-append window, and two tasks racing (even both past an expired seed lock) can
+   * never both write a seed. Returns true if THIS call wrote the seed (apply it locally), false if the
+   * stream was already seeded (the tailer will deliver the peer's seed — do NOT apply a second one).
+   * Retries a transient Redis error like {@link appendUpdate}; throws if it ultimately fails. Disabled →
+   * true (single-replica: seed locally, no stream).
+   */
+  async seedIfEmpty(name: string, update: Uint8Array): Promise<boolean> {
+    if (!this.enabled || !this.write) return true
+    const encoded = Buffer.from(update).toString('base64')
+    for (let attempt = 0; attempt <= PUBLISH_MAX_RETRIES; attempt++) {
+      try {
+        const wrote = await this.write.eval(SEED_IF_EMPTY_SCRIPT, {
+          keys: [streamKey(name)],
+          arguments: [UPDATE_FIELD, encoded],
+        })
+        await this.write.expire(streamKey(name), STREAM_TTL_SEC).catch(() => {})
+        return wrote === 1
+      } catch (error) {
+        if (attempt === PUBLISH_MAX_RETRIES) {
+          logger.error(`FileDocStore seed failed for ${name}`, { error: getErrorMessage(error) })
+          throw error
+        }
+        await sleep(backoffWithJitter(attempt + 1, null, { baseMs: 50, maxMs: 500 }))
+      }
+    }
+    return false
+  }
+
+  /**
+   * Whether the file's stream already holds content — an EFFICIENCY recheck in {@link shouldSeed} that
+   * skips the seed fetch when a prior holder already seeded (the split-brain guard itself is the atomic
+   * {@link SEED_IF_EMPTY_SCRIPT}, not this check). Treats `true` as "already seeded", so it fails CLOSED:
+   * a Redis `xLen` error returns `true` (cannot confirm empty → skip the redundant fetch; the atomic seed
+   * would no-op anyway). `false` only when genuinely empty, or when disabled (single-replica).
    */
   async streamHasContent(name: string): Promise<boolean> {
     if (!this.enabled || !this.write) return false
@@ -321,17 +366,20 @@ export class FileDocStore {
   }
 
   /**
-   * Decide whether THIS task should build and write the file's one-time seed. Returns a lock TOKEN only
-   * when the shared stream is genuinely empty AND this task wins the seed lock — so exactly one task
-   * across the cluster ever seeds a file, even if several open it at once (the fix for split-brain
-   * seeding). `null` otherwise. Release the token with {@link releaseSeedLock}. Disabled → always a
-   * token (single-replica: seed locally).
+   * Decide whether THIS task should run the (expensive) seed fetch + write for a file. Returns a lock
+   * TOKEN when this task wins the seed lock and the stream still looks empty; `null` otherwise. This is an
+   * EFFICIENCY gate — it stops every task that opens the file at once from each fetching the seed. It does
+   * NOT by itself guarantee a single seed: exactly-once is enforced by the atomic {@link seedIfEmpty} the
+   * token-holder then calls (the split-brain guard), so a lock that expires mid-seed cannot cause a
+   * double-seed. Release the token with {@link releaseSeedLock}. Disabled → always a token (single-replica:
+   * seed locally).
    */
   async shouldSeed(name: string): Promise<string | null> {
     const token = await this.acquireLock(`${SEED_LOCK_PREFIX}${name}`, SEED_LOCK_TTL_MS)
     if (!token || token === DISABLED_LOCK_TOKEN) return token
     // The lock could be free yet the stream already seeded (a prior holder seeded then its lock
-    // expired). Re-check under the lock so we never write a SECOND seed on top of an existing one.
+    // expired). Re-check so we skip the redundant seed fetch — the atomic seedIfEmpty would no-op anyway,
+    // but this avoids the wasted app round-trip.
     if (await this.streamHasContent(name)) {
       await this.releaseSeedLock(name, token)
       return null

apps/realtime/src/handlers/file-doc.ts

@@ -373,24 +373,25 @@ async function ensureServerSeed(
   try {
     const update = await fetchFileDocSeed(workspaceId, room.fileId)
     if (fileDocRooms.get(name) !== room || isDocSeeded(room.doc)) return
-    // Fence against a peer that seeded while we held a stale lock (a rare long stall): if the stream
-    // already has content it is seeded, so never write a SECOND seed on top — that would split-brain.
-    if (await store.streamHasContent(name)) {
-      // Clear the guard so a later join can retry. Safe in both cases: a genuine peer-seed arrives via
-      // the tailer (and shouldSeed/this fence then no-op), and a fail-closed Redis `xLen` error must not
-      // strand the room unseeded forever.
-      room.serverSeedStarted = false
-      return
-    }
-    // Build the seed (file content + seed flag, or just the flag for an empty/missing file), then
-    // PUBLISH it to the shared stream AWAITED *before* seeding the local doc. The doc is marked seeded
-    // only once the seed is durably shared — so if the publish fails we leave the doc unseeded and the
-    // stream empty for a clean retry, rather than serving an unpublished local seed that a peer would
-    // re-seed on top of (split-brain). SEED_ORIGIN keeps `doc.on('update')` from re-publishing it.
+    // Build the seed (file content + seed flag, or just the flag for an empty/missing file) and write it
+    // to the shared stream ATOMICALLY, iff the stream is still empty. This — NOT the seed lock — is the
+    // split-brain guard: two tasks racing (even both past an expired lock) can never both seed, because
+    // the emptiness check and the append are one Redis-side step. Publish-before-apply: the doc is marked
+    // seeded (via the local apply) only once the seed is durably in the stream, so a failed write leaves
+    // the doc unseeded and the stream empty for a clean retry. SEED_ORIGIN keeps `doc.on('update')` from
+    // re-publishing it.
     const seedUpdate = update ?? emptySeedUpdate()
-    await store.publishAndWait(name, seedUpdate)
+    const didSeed = await store.seedIfEmpty(name, seedUpdate)
     if (fileDocRooms.get(name) !== room || isDocSeeded(room.doc)) return
-    Y.applyUpdate(room.doc, seedUpdate, SEED_ORIGIN)
+    if (didSeed) {
+      Y.applyUpdate(room.doc, seedUpdate, SEED_ORIGIN)
+    } else {
+      // A peer seeded first: its seed arrives via the tailer, so we must NOT apply our own — a second,
+      // different-client-id seed IS the split-brain. Clear the guard so a later join can retry if that
+      // peer seed somehow never lands (e.g. a fail-closed `xLen` error made `shouldSeed` skip a genuinely
+      // empty stream); a real peer-seed makes the retry a no-op.
+      room.serverSeedStarted = false
+    }
   } catch (error) {
     logger.warn(`Server seed failed for file ${room.fileId} (workspace ${workspaceId})`, error)
     room.serverSeedStarted = false
@@ -549,8 +550,11 @@ function getOrCreateRoom(io: Server, ref: RoomRef): FileDocRoom {
     // whichever task is last to leave persists real edits, even one that only tailed them. A compaction
     // snapshot on catch-up (REDIS_SNAPSHOT_ORIGIN) also counts: it folds real edits into one frame, so a
     // fresh task catching up purely from it must not treat the doc as unedited. The seed transition
-    // itself is never counted (nor a purely-local copilot merge, already durable via copilot's direct
-    // file write), so a seeded-but-unedited doc is never projected back over the file.
+    // itself is never counted, so a seeded-but-unedited doc is never projected back over the file. NOTE:
+    // in the multi-replica path a copilot merge is NOT excluded — it round-trips through the stream as
+    // REDIS_ORIGIN, indistinguishable from a peer edit, so it marks `edited`. That only ever causes an
+    // extra idempotent persist of content copilot already wrote directly (safe over-persist, never a lost
+    // edit); the single-replica path applies the merge locally with no origin and does not count it.
     const seededBefore = room.seededObserved
     if (isDocSeeded(room.doc)) room.seededObserved = true
     if (