PotentialStorage/test/kv-storage/kv-storage-composite.test.ts

import { afterEach, describe, expect, it, vi } from 'vitest'
import { KVStorageComposite } from '../../src/kv-storage/kv-storage-composite.js'
import { KVStorageInMemory } from '../../src/kv-storage/kv-storage-in-memory.js'
import type { BaseKVStorage, KVChangeEvent } from '../../src/kv-storage/base-kv-storage.js'
import { DocumentStorageKV } from '../../src/document-storage/document-storage-kv.js'
import { collect, KVStorageCounting } from '../helpers.js'
import { runKVStorageConformance } from './kv-storage.conformance.js'
import { TestNetwork } from './sync-helpers.js'

runKVStorageConformance(
  'KVStorageComposite (in-memory source, 1 replica)',
  () => new KVStorageComposite({ source: new KVStorageInMemory(), replicas: [new KVStorageInMemory()] }),
)
runKVStorageConformance('KVStorageComposite (http source, 2 replicas)', () => {
  network = new TestNetwork()
  return new KVStorageComposite({ source: network.client(), replicas: [new KVStorageInMemory(), new KVStorageInMemory()] })
})

let network: TestNetwork
const composites: KVStorageComposite[] = []
afterEach(async () => {
  await Promise.all(composites.splice(0).map((c) => c.close()))
  await network?.close()
})

function composite(persistent: BaseKVStorage = new KVStorageInMemory(), queue?: BaseKVStorage) {
  const c = new KVStorageComposite({ source: network.client(), replicas: [new KVStorageInMemory(), persistent], queue, retryMs: 10 })
  composites.push(c)
  return c
}

describe('KVStorageComposite sync', () => {
  it('loads existing data with one sync, then serves reads locally', async () => {
    network = new TestNetwork()
    await network.store.writeMany([
      { key: 'users/1', value: { name: 'Ada' } },
      { key: 'users/2', value: { name: 'Bob' } },
    ])
    const db = composite()
    await db.sync()
    network.store.reset()
    expect(await db.child('users').get('1')).toEqual({ name: 'Ada' })
    expect(await collect(db.child('users').keys())).toEqual(['1', '2'])
    expect(network.store.calls.readMany).toBeUndefined() // (a follow-up sync started on subscribe may still scan)
  })

  it('applies changes pushed from other clients to every replica', async () => {
    network = new TestNetwork()
    const memory = new KVStorageInMemory()
    const db = composite(memory)
    await db.sync()
    const seen = db.waitFor('change', (e) => e.origin === 'remote')
    await network.client().child('users').set('9', { name: 'Remote' })
    expect((await seen).changes).toEqual([{ key: 'users/9', value: { name: 'Remote' } }])
    expect(await db.child('users').get('9')).toEqual({ name: 'Remote' })
    expect(await memory.readMany(['users/9'])).toEqual([{ name: 'Remote' }])
  })

  it('announces only remote changes that differ from local values', async () => {
    network = new TestNetwork()
    const db = composite()
    await db.sync()
    const events: KVChangeEvent[] = []
    db.on('change', (e) => events.push(e))
    await db.set('mine', 1) // relayed back by the server: same value, so no remote event
    const other = db.waitFor('change', (e) => e.origin === 'remote')
    await network.client().writeMany([
      { key: 'mine', value: 1 },
      { key: 'theirs', value: 2 },
    ])
    await other
    expect(events).toEqual([
      { changes: [{ key: 'mine', value: 1 }], origin: 'local' },
      { changes: [{ key: 'theirs', value: 2 }], origin: 'remote' },
    ])
  })

  it('queues writes while offline and replays them in order when back online', async () => {
    network = new TestNetwork()
    const db = composite()
    await db.sync()
    network.online = false
    await db.child('users').set('1', { v: 1 })
    await db.child('users').set('1', { v: 2 })
    await db.child('users').delete('1')
    await db.child('users').set('2', { v: 3 })
    expect(db.queued).toBe(4)
    expect(await db.child('users').get('2')).toEqual({ v: 3 }) // served locally
    network.online = true
    await vi.waitFor(() => expect(db.queued).toBe(0))
    expect(await network.store.readMany(['users/1', 'users/2'])).toEqual([undefined, { v: 3 }])
  })

  it('the last write wins when the server pushes a change while a local write is queued', async () => {
    network = new TestNetwork()
    const db = composite()
    await db.sync()
    network.online = false
    await db.set('k', { from: 'local' })
    await network.store.writeMany([{ key: 'k', value: { from: 'remote' } }]) // pushed over the still-open stream
    await vi.waitFor(async () => expect(await db.get('k')).toEqual({ from: 'remote' }))
    network.online = true
    await vi.waitFor(() => expect(db.queued).toBe(0)) // the queued write is sent last, so it wins
    expect(await network.store.readMany(['k'])).toEqual([{ from: 'local' }])
    expect(await db.get('k')).toEqual({ from: 'local' })
  })

  it('restarts from the persistent replica, then catches up', async () => {
    network = new TestNetwork()
    const persistent = new KVStorageInMemory()
    const queue = new KVStorageInMemory()
    const first = composite(persistent, queue)
    await first.set('a', 1)
    await first.set('gone', 1)
    await first.sync()
    await first.close()
    await network.client().writeMany([{ key: 'b', value: 2 }, { key: 'gone' }]) // missed while closed

    network.online = false
    const second = composite(persistent, queue)
    expect(await second.get('a')).toBe(1) // available without the server
    network.online = true
    await second.sync()
    expect(await collect(second.scan({ values: true }))).toEqual([
      ['a', 1],
      ['b', 2],
    ])
  })

  it('replays writes queued by a previous session from a persistent queue', async () => {
    network = new TestNetwork()
    const persistent = new KVStorageInMemory()
    const queue = persistent.child('~sync')
    const first = new KVStorageComposite({ source: network.client(), replicas: [persistent], queue, prefix: 'data' })
    composites.push(first)
    await first.sync()
    network.online = false
    await first.child('data').set('offline', 'queued')
    await first.close()
    network.online = true
    const second = new KVStorageComposite({ source: network.client(), replicas: [persistent], queue, prefix: 'data' })
    composites.push(second)
    expect(second.queued).toBe(0)
    await second.sync()
    expect(second.queued).toBe(0)
    expect(await network.store.readMany(['data/offline'])).toEqual(['queued'])
  })

  it('catches up on changes missed while the stream was down', async () => {
    network = new TestNetwork()
    const db = composite()
    await db.sync()
    const resynced = vi.fn()
    db.on('change', resynced)
    network.online = false
    network.dropStreams()
    await network.store.writeMany([{ key: 'missed', value: 1 }])
    network.online = true
    await vi.waitFor(async () => expect(await db.get('missed')).toBe(1), { timeout: 5000 })
    expect(resynced).toHaveBeenCalledWith({ changes: [{ key: 'missed', value: 1 }], origin: 'remote' })
  })

  it('syncs only its prefix and rejects keys outside it', async () => {
    network = new TestNetwork()
    await network.store.writeMany([
      { key: 'alice/a', value: 1 },
      { key: 'bob/b', value: 2 },
    ])
    const db = new KVStorageComposite({ source: network.client(), replicas: [new KVStorageInMemory()], prefix: 'alice' })
    composites.push(db)
    await db.sync()
    expect(await collect(db.scan({ values: true }))).toEqual([['alice/a', 1]])
    await expect(db.child('bob').get('b')).rejects.toBeInstanceOf(RangeError)
  })

  it('refuses a queue inside the synced range of a replica', () => {
    const replica = new KVStorageInMemory()
    expect(() => new KVStorageComposite({ source: new KVStorageInMemory(), replicas: [replica], queue: replica.child('q') })).toThrow(/outside/)
    expect(() => new KVStorageComposite({ source: new KVStorageInMemory(), replicas: [replica], queue: replica.child('q'), prefix: 'data' })).not.toThrow()
  })

  it('serves document queries from the local replica, offline too', async () => {
    network = new TestNetwork()
    const db = composite()
    const users = new DocumentStorageKV(db).child('users')
    await users.set('1', { age: 30 })
    await users.set('2', { age: 40 })
    await db.sync()
    network.online = false
    expect((await collect(users.find({ age: { $gt: 35 } }))).map((r) => r.key)).toEqual(['2'])
  })

  it('does not undo writes made while a sync is scanning', async () => {
    // A source whose scan reads ahead (like the HTTP stream) and pauses after its first entry.
    let release!: () => void
    let reached!: () => void
    const gate = new Promise<void>((r) => (release = r))
    const paused = new Promise<void>((r) => (reached = r))
    const inner = new KVStorageInMemory()
    await inner.writeMany([
      { key: 'a', value: 1 },
      { key: 'c', value: 3 },
    ])
    const source = new KVStorageCounting(inner)
    const scan = source.scan.bind(source)
    source.scan = async function* (options) {
      let first = true
      for (const entry of await collect(scan(options))) {
        yield entry
        if (first && options?.values) ((first = false), reached(), await gate)
      }
    }
    const persistent = new KVStorageInMemory()
    await persistent.writeMany([{ key: 'a0', value: 'stale' }])
    const db = new KVStorageComposite({ source, replicas: [persistent] })
    composites.push(db)
    void db.start()
    await paused // the diff has passed 'b' in the source but not in the replica
    await db.set('b', 2)
    release()
    await vi.waitFor(() => expect(db.queued).toBe(0))
    await db.sync()
    expect(await collect(db.scan({ values: true }))).toEqual([
      ['a', 1],
      ['b', 2],
      ['c', 3],
    ])
  })

  it('sync() during a running sync runs one more, so changes after its scan are picked up', async () => {
    let release!: () => void
    let reached!: () => void
    const gate = new Promise<void>((r) => (release = r))
    const paused = new Promise<void>((r) => (reached = r))
    const inner = new KVStorageInMemory()
    await inner.writeMany([{ key: 'a', value: 1 }])
    const source = new KVStorageCounting(inner)
    const scan = source.scan.bind(source)
    let scans = 0
    source.scan = async function* (options) {
      const entries = await collect(scan(options)) // a snapshot, like a read-ahead stream
      if (options?.values && scans++ === 0) (reached(), await gate)
      yield* entries
    }
    const db = new KVStorageComposite({ source, replicas: [new KVStorageInMemory()] })
    composites.push(db)
    const first = db.sync()
    await paused
    await inner.writeMany([{ key: 'b', value: 2 }]) // after the running sync's snapshot
    const second = db.sync()
    expect(db.sync()).toBe(second) // later callers share the pending one
    release()
    await first
    await second
    expect(await collect(db.scan({ values: true }))).toEqual([
      ['a', 1],
      ['b', 2],
    ])
  })
})