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

import { StorageUnavailableError } from '../utils/errors.js'
import type { OffCallback } from '../utils/event-emitter.js'
import { Json } from '../utils/json.js'
import { KeyRange } from '../utils/key-range.js'
import type { AnySpec, ChildOf, KeyOf, Spec } from '../utils/schema.js'
import { Ulid } from '../utils/ulid.js'
import { BaseKVStorage, type KVChange, type ScanOptions } from './base-kv-storage.js'
import { KVStorageInMemory } from './kv-storage-in-memory.js'

export interface KVStorageCompositeOptions {
  /** The source of truth, e.g. `KVStorageHttp`. */
  source: BaseKVStorage
  /** Local copies, fastest first, e.g. `[inMemory, indexedDb]`. The last one seeds the others on start. */
  replicas: BaseKVStorage[]
  /**
   * Holds writes made while the source is unreachable (under its `writes` child) and whether a
   * sync has completed (`synced`). Use a persistent node outside `prefix`, e.g.
   * `indexedDb.child('sync')`, so queued writes survive a reload. Default: in memory.
   */
  queue?: BaseKVStorage
  /** Only this key and its subtree are synced (e.g. the user's public key). Default: everything. */
  prefix?: string
  /** Delay before retrying when the source is unreachable. Default 5 000 ms. */
  retryMs?: number
}

/** Sources that push remote changes (e.g. `KVStorageHttp`). */
interface Subscribable {
  subscribe(path: string): Promise<void>
  unsubscribe(path: string): Promise<void>
}

// Types only: children of a KVStorageComposite are KVStorageComposite nodes too (child() creates
// them with Object.create, so they share this class's prototype).
export interface KVStorageComposite<S extends Spec = AnySpec> {
  child<K extends KeyOf<S>>(key: K): KVStorageComposite<ChildOf<S, K>>
}

/**
 * Combines a source with local replicas: reads come from the first replica, writes go to the
 * source and then every replica, exactly as given (no merging: the last write wins).
 *
 * - Changes others make arrive from the source (`subscribe`) and are applied to the replicas.
 * - A write the source can't take (`StorageUnavailableError`) is applied locally and stored in
 *   the `queue` store; later writes queue behind it, so order is kept.
 * - `sync()` replays the queue, then copies what changed on the source into the replicas. It runs
 *   on start, when the source reconnects (`resync`), and after a retry delay while offline.
 * - Before the first successful sync, reads go to the source.
 */
export class KVStorageComposite<S extends Spec = AnySpec> extends BaseKVStorage<S> {
  private readonly source: BaseKVStorage
  private readonly replicas: BaseKVStorage[]
  private readonly queue: BaseKVStorage
  private readonly writes: BaseKVStorage
  private readonly range: KeyRange
  private readonly retryMs: number
  private readonly ids = new Ulid()
  /** Mutable state lives here: child nodes share this object (see `child()`). */
  private readonly state: {
    started?: Promise<void>
    syncing?: Promise<void>
    /** A sync has completed at some point, so the replicas hold a full copy. */
    synced: boolean
    /** Writes waiting in the queue. */
    queued: number
    /** Remote changes are applied one after another. */
    applying: Promise<unknown>
    retry?: ReturnType<typeof setTimeout>
    offs: OffCallback[]
  } = { synced: false, queued: 0, applying: Promise.resolve(), offs: [] }

  constructor(options: KVStorageCompositeOptions) {
    super()
    if (options.replicas.length === 0) throw new TypeError('KVStorageComposite needs at least one replica')
    this.source = options.source
    this.replicas = options.replicas
    this.queue = options.queue ?? new KVStorageInMemory()
    this.writes = this.queue.child('writes')
    this.range = new KeyRange(options.prefix ?? '')
    this.retryMs = options.retryMs ?? 5000
  }

  /** Loads local data and starts syncing (also happens on first use). */
  start(): Promise<void> {
    return (this.state.started ??= this.open())
  }

  /** Replays queued writes, then brings the replicas up to date with the source. */
  sync(): Promise<void> {
    return (this.state.syncing ??= this.runSync().finally(() => (this.state.syncing = undefined)))
  }

  /** Number of writes waiting for the source. */
  get queued(): number {
    return this.state.queued
  }

  async close(): Promise<void> {
    clearTimeout(this.state.retry)
    for (const off of this.state.offs.splice(0)) off()
    if (KVStorageComposite.isSubscribable(this.source)) await this.source.unsubscribe(this.range.prefix)
  }

  async readMany(keys: string[]): Promise<unknown[]> {
    this.checkRange(keys)
    await this.start()
    return (this.state.synced ? this.replicas[0]! : this.source).readMany(keys)
  }

  async *scan(options: ScanOptions = {}): AsyncGenerator<[string, unknown]> {
    await this.start()
    const scoped = this.range.clamp(options)
    if (scoped) yield* (this.state.synced ? this.replicas[0]! : this.source).scan(scoped)
  }

  protected async write(changes: KVChange[]): Promise<void> {
    this.checkRange(changes.map(({ key }) => key))
    await this.start()
    // While syncing or with writes queued, queue this one too so the source sees writes in order.
    if (this.state.queued === 0 && !this.state.syncing) {
      try {
        await this.source.writeMany(changes)
        return this.apply(changes)
      } catch (error) {
        if (!(error instanceof StorageUnavailableError)) throw error
      }
    }
    await this.writes.set(this.ids.generate(), changes)
    this.state.queued++
    await this.apply(changes)
    this.retryLater()
  }

  private async open(): Promise<void> {
    this.state.synced = (await this.queue.get('synced')) === true
    for await (const _ of this.writes.keys()) this.state.queued++
    await this.seed()
    this.state.offs.push(
      this.source.on('change', (event) => {
        if (event.origin !== 'remote') return
        const changes = event.changes.filter(({ key }) => this.range.contains(key))
        this.state.applying = this.state.applying.then(() => this.apply(changes, true)).catch((error) => console.error(error))
      }),
      this.source.on('resync', () => void this.sync()),
    )
    if (KVStorageComposite.isSubscribable(this.source)) {
      void this.source
        .subscribe(this.range.prefix)
        .then(() => this.sync())
        .catch((error) => console.error(error))
    } else {
      void this.sync()
    }
  }

  private async runSync(): Promise<void> {
    try {
      await this.replay()
      let batch: KVChange[] = []
      for await (const change of this.diff()) if (batch.push(change) >= 500) await this.apply(batch.splice(0), true)
      await this.apply(batch, true)
      await this.replay() // writes made during the diff: send them, and re-apply them over it
      await this.queue.set('synced', true)
      this.state.synced = true
    } catch (error) {
      if (!(error instanceof StorageUnavailableError)) console.error(error)
      this.retryLater()
    }
  }

  /** Sends queued writes to the source in order and re-applies them locally. */
  private async replay(): Promise<void> {
    for (;;) {
      let next: [string, unknown] | undefined
      for await (const entry of this.writes.entries({ limit: 1 })) next = entry
      if (!next) return
      const changes = next[1] as KVChange[]
      await this.source.writeMany(changes)
      await this.apply(changes)
      await this.writes.delete(next[0])
      this.state.queued--
    }
  }

  /** What the source has that differs from the last replica (both scanned in key order). */
  private async *diff(): AsyncGenerator<KVChange> {
    const options = { gte: this.range.gte, lt: this.range.lt, values: true }
    const remote = this.source.scan(options)[Symbol.asyncIterator]()
    const local = this.replicas.at(-1)!.scan(options)[Symbol.asyncIterator]()
    let r = await remote.next()
    let l = await local.next()
    while (!r.done || !l.done) {
      if (l.done || (!r.done && r.value[0] < l.value[0])) {
        yield { key: r.value![0], value: r.value![1] }
        r = await remote.next()
      } else if (r.done || l.value[0] < r.value[0]) {
        yield { key: l.value[0] }
        l = await local.next()
      } else {
        if (!Json.equal(r.value[1], l.value[1])) yield { key: r.value[0], value: r.value[1] }
        ;[r, l] = await Promise.all([remote.next(), local.next()])
      }
    }
  }

  /** Writes changes to every replica; `remote` ones are also announced as `change` events. */
  private async apply(changes: KVChange[], remote = false): Promise<void> {
    if (changes.length === 0) return
    await Promise.all(this.replicas.map((replica) => replica.writeMany(changes)))
    if (remote) this.emit('change', { changes, origin: 'remote' })
  }

  /** Copies the last (persistent) replica into the faster ones. */
  private async seed(): Promise<void> {
    if (this.replicas.length < 2) return
    let batch: KVChange[] = []
    const flush = () => Promise.all(this.replicas.slice(0, -1).map((replica) => replica.writeMany(batch.splice(0))))
    for await (const [key, value] of this.replicas.at(-1)!.scan({ gte: this.range.gte, lt: this.range.lt, values: true })) {
      if (batch.push({ key, value }) >= 500) await flush()
    }
    await flush()
  }

  private retryLater(): void {
    clearTimeout(this.state.retry)
    this.state.retry = setTimeout(() => void this.sync(), this.retryMs)
    ;(this.state.retry as { unref?: () => void }).unref?.()
  }

  private checkRange(keys: string[]): void {
    for (const key of keys) if (!this.range.contains(key)) throw new RangeError(`"${key}" is outside this composite's prefix "${this.range.prefix}"`)
  }

  private static isSubscribable(source: BaseKVStorage): source is BaseKVStorage & Subscribable {
    return typeof (source as Partial<Subscribable>).subscribe === 'function'
  }
}