dorfteich/apps/collab/src/persistence.ts

import { MAX_PAGE_DOCUMENT_BYTES } from '@dorfteich/shared';
import type { Pool } from 'pg';
import * as Y from 'yjs';

import { deriveContentFromDoc } from './yjs-content.js';

/** Outcome of a store attempt, surfaced to the caller for logging/notification. */
export interface StoreResult {
  outcome: 'stored' | 'too_large' | 'not_found';
  /** Size of the full merged state in bytes (for logging and the size ceiling). */
  bytes: number;
  /** Wall-clock duration of the flush, in milliseconds (ADR: measure flushes). */
  durationMs: number;
  /** Whether this flush merged the update log back into `pages.ydoc_state`. */
  merged: boolean;
}

/**
 * The persistence port the collab hooks depend on (ADR 0003). Kept as an
 * interface so the server hooks can be integration-tested with an in-memory
 * fake, while the Postgres implementation is exercised by a DB-backed test.
 */
export interface PagePersistence {
  /**
   * Apply the persisted state of `pageId` into `doc`. Returns `false` when the
   * page does not exist or is trashed — the caller then keeps the empty doc.
   */
  loadInto(pageId: string, doc: Y.Doc): Promise<boolean>;
  /** Persist the current `doc`, enforcing the document size ceiling. */
  store(pageId: string, doc: Y.Doc): Promise<StoreResult>;
  /** Release per-document bookkeeping when Hocuspocus unloads the document. */
  forget(pageId: string): void;
}

/**
 * When the append log for an open document grows past this many rows, a store
 * flush merges it back into `pages.ydoc_state` and truncates the log. This
 * bounds the cost of `loadInto` for long-lived sessions; the heavier,
 * session-aware compaction of idle pages is the separate maintenance job (#40,
 * default threshold 500 in realtime-collaboration.md).
 */
const INLINE_MERGE_THRESHOLD = 200;

export class PostgresPagePersistence implements PagePersistence {
  /**
   * The Yjs state vector last persisted for each open document, so each store
   * appends only the delta since the previous flush. Seeded by `loadInto` and
   * cleared by `forget`; a missing entry safely falls back to storing the full
   * state as the delta.
   */
  private readonly lastStoredVector = new Map<string, Uint8Array>();

  constructor(private readonly pool: Pool) {}

  async loadInto(pageId: string, doc: Y.Doc): Promise<boolean> {
    const pageRow = await this.pool.query<{ ydoc_state: Buffer }>(
      'SELECT ydoc_state FROM pages WHERE id = $1 AND deleted_at IS NULL',
      [pageId],
    );
    const stored = pageRow.rows[0];
    if (!stored) return false;

    Y.applyUpdate(doc, new Uint8Array(stored.ydoc_state));

    const updates = await this.pool.query<{ update: Buffer }>(
      'SELECT update FROM page_updates WHERE page_id = $1 ORDER BY seq ASC',
      [pageId],
    );
    for (const row of updates.rows) {
      Y.applyUpdate(doc, new Uint8Array(row.update));
    }

    this.lastStoredVector.set(pageId, Y.encodeStateVector(doc));
    return true;
  }

  async store(pageId: string, doc: Y.Doc): Promise<StoreResult> {
    const start = performance.now();
    // Capture everything from the live doc synchronously, before any await, so
    // concurrent inbound updates cannot change what this flush persists.
    const full = Y.encodeStateAsUpdate(doc);
    const durationOf = (): number => performance.now() - start;

    if (full.byteLength > MAX_PAGE_DOCUMENT_BYTES) {
      return {
        outcome: 'too_large',
        bytes: full.byteLength,
        durationMs: durationOf(),
        merged: false,
      };
    }

    const previousVector = this.lastStoredVector.get(pageId);
    const delta = previousVector ? Y.encodeStateAsUpdate(doc, previousVector) : full;
    const derived = deriveContentFromDoc(doc);
    const nextVector = Y.encodeStateVector(doc);

    const client = await this.pool.connect();
    try {
      await client.query('BEGIN');

      // Lock the page row for the duration of the flush: this serialises seq
      // allocation and guards against storing to a page trashed mid-session.
      const page = await client.query<{ pond_id: string }>(
        'SELECT pond_id FROM pages WHERE id = $1 AND deleted_at IS NULL FOR UPDATE',
        [pageId],
      );
      const pageMeta = page.rows[0];
      if (!pageMeta) {
        await client.query('ROLLBACK');
        return {
          outcome: 'not_found',
          bytes: full.byteLength,
          durationMs: durationOf(),
          merged: false,
        };
      }
      const pondId = pageMeta.pond_id;

      const seqRow = await client.query<{ seq: number }>(
        'SELECT COALESCE(MAX(seq) + 1, 0) AS seq FROM page_updates WHERE page_id = $1',
        [pageId],
      );
      const seq = seqRow.rows[0]?.seq ?? 0;
      await client.query(
        'INSERT INTO page_updates (id, page_id, seq, update) VALUES (gen_random_uuid(), $1, $2, $3)',
        [pageId, seq, Buffer.from(delta)],
      );

      const merged = seq + 1 >= INLINE_MERGE_THRESHOLD;
      if (merged) {
        await client.query('UPDATE pages SET ydoc_state = $2, updated_at = now() WHERE id = $1', [
          pageId,
          Buffer.from(full),
        ]);
        await client.query('DELETE FROM page_updates WHERE page_id = $1', [pageId]);
      } else {
        await client.query('UPDATE pages SET updated_at = now() WHERE id = $1', [pageId]);
      }

      await client.query(
        `INSERT INTO page_content_cache (page_id, plain_text, markdown, html, outline, updated_at)
         VALUES ($1, $2, $3, $4, $5::jsonb, now())
         ON CONFLICT (page_id) DO UPDATE
           SET plain_text = EXCLUDED.plain_text,
               markdown = EXCLUDED.markdown,
               html = EXCLUDED.html,
               outline = EXCLUDED.outline,
               updated_at = now()`,
        [
          pageId,
          derived.plainText,
          derived.markdown,
          derived.html,
          JSON.stringify(derived.outline),
        ],
      );

      if (derived.imageFileIds.length > 0) {
        // Keep Attachment.pageId pointed at the page embedding the file (#31),
        // scoped to the pond so a client cannot claim another pond's file.
        await client.query(
          'UPDATE attachments SET page_id = $1 WHERE id = ANY($2::text[]) AND pond_id = $3',
          [pageId, derived.imageFileIds, pondId],
        );
      }

      await client.query('COMMIT');
      this.lastStoredVector.set(pageId, nextVector);
      return { outcome: 'stored', bytes: full.byteLength, durationMs: durationOf(), merged };
    } catch (error) {
      await client.query('ROLLBACK').catch(() => undefined);
      throw error;
    } finally {
      client.release();
    }
  }

  forget(pageId: string): void {
    this.lastStoredVector.delete(pageId);
  }
}