feat(localdb): snapshot resync (LWW-merging, tombstone-carrying, delta resume)

Claude-Session: https://claude.ai/code/session_01BL2Vu1ESDQ9SCN4gVKkdts
2026-07-17 23:35:50 -04:00
parent e92cdc6d3a
commit 130e9c918f
5 changed files with 628 additions and 27 deletions
@@ -0,0 +1,171 @@
+using Microsoft.Extensions.Logging;
+using ZB.MOM.WW.LocalDb.Contracts;
+using ZB.MOM.WW.LocalDb.Internal;
+using ZB.MOM.WW.LocalDb.Registration;
+
+namespace ZB.MOM.WW.LocalDb.Replication.Internal;
+
+/// <summary>
+/// Consumes an inbound full-snapshot the receive loop routes to it: SnapshotBegin opens a receive,
+/// each SnapshotBatch merges through the SAME last-writer-wins path as deltas (never truncating —
+/// the receiver may hold newer local writes), and SnapshotComplete advances the durable watermark
+/// so incremental deltas resume above it.
+/// </summary>
+internal interface ISnapshotApplier
+{
+    Task OnBeginAsync(SnapshotBegin begin, CancellationToken ct);
+    Task OnBatchAsync(SnapshotBatch batch, CancellationToken ct);
+    Task OnCompleteAsync(SnapshotComplete complete, CancellationToken ct);
+}
+
+/// <summary>
+/// Streams a full snapshot of every registered table to a peer that has fallen behind the pruned
+/// oplog horizon. Reads the authoritative per-row version from <c>__localdb_row_version</c> joined
+/// to each live table, pages it into <see cref="SnapshotBatch"/> messages through the gated data
+/// lane, and brackets them with SnapshotBegin/SnapshotComplete carrying the max oplog seq as-of.
+/// </summary>
+internal sealed class SnapshotStreamer(SqliteLocalDb db, OplogStore store, ReplicationOptions options, ILogger logger)
+{
+    public async Task SendAsync(Func<SyncMessage, CancellationToken, Task> send, CancellationToken ct)
+    {
+        var asOfSeq = await store.GetMaxSeqAsync(ct);
+        await send(new SyncMessage { SnapshotBegin = new SnapshotBegin { AsOfSeq = asOfSeq } }, ct);
+
+        var totalRows = 0L;
+        foreach (var table in db.ReplicatedTables.Values.OrderBy(t => t.Name, StringComparer.Ordinal))
+        {
+            var sql = BuildPageSql(table);
+            var offset = 0;
+            while (true)
+            {
+                var rows = await db.QueryAsync(
+                    sql,
+                    // Column 4 is NULL for tombstones; a proto optional string cannot be assigned null,
+                    // so carry it as a nullable and only set row_json when present.
+                    static r => (
+                        PkJson: r.GetString(0),
+                        Hlc: r.GetInt64(1),
+                        NodeId: r.GetString(2),
+                        IsTombstone: r.GetInt64(3) != 0,
+                        RowJson: r.IsDBNull(4) ? null : r.GetString(4)),
+                    new { t = table.Name, page = options.MaxBatchSize, off = offset }, ct);
+                if (rows.Count == 0)
+                    break;
+
+                var batch = new SnapshotBatch { TableName = table.Name };
+                foreach (var row in rows)
+                {
+                    var snapshotRow = new SnapshotRow
+                    {
+                        PkJson = row.PkJson,
+                        Hlc = row.Hlc,
+                        NodeId = row.NodeId,
+                        IsTombstone = row.IsTombstone,
+                    };
+                    if (row.RowJson is not null)
+                        snapshotRow.RowJson = row.RowJson;
+                    batch.Rows.Add(snapshotRow);
+                }
+                await send(new SyncMessage { SnapshotBatch = batch }, ct);
+                totalRows += rows.Count;
+
+                offset += rows.Count;
+                if (rows.Count < options.MaxBatchSize)
+                    break;
+            }
+        }
+
+        await send(new SyncMessage { SnapshotComplete = new SnapshotComplete { AsOfSeq = asOfSeq } }, ct);
+        // The sender owed this snapshot (its own needs_snapshot); clearing it here — after the whole
+        // snapshot is on the wire — stops it re-streaming on the next reconnect.
+        await store.SetNeedsSnapshotAsync(false, ct);
+        logger.LogInformation("Snapshot sent (as_of_seq {AsOfSeq}, {Rows} rows).", asOfSeq, totalRows);
+    }
+
+    // Pages one table's rows by joining the row-version ledger to the live table on the PK columns
+    // extracted from pk_json. Rows present in the live table but absent from __localdb_row_version
+    // (data written BEFORE the table was registered for replication) are intentionally NOT snapshotted:
+    // they were never captured, consistent with CDC semantics — a full re-registration is required to
+    // baseline them.
+    private static string BuildPageSql(ReplicatedTable table)
+    {
+        var name = Ident(table.Name);
+        var join = string.Join(" AND ", table.PkColumns.Select(p => $"t.{Ident(p)} = json_extract(rv.pk_json, {JsonPath(p)})"));
+        var jsonObject = "json_object(" +
+            string.Join(", ", table.Columns.Select(c => $"{StringLiteral(c.Name)}, t.{Ident(c.Name)}")) + ")";
+        return
+            "SELECT rv.pk_json, rv.hlc, rv.node_id, rv.is_tombstone, " +
+            $"CASE WHEN rv.is_tombstone = 1 THEN NULL ELSE {jsonObject} END " +
+            $"FROM __localdb_row_version rv LEFT JOIN {name} t ON {join} " +
+            "WHERE rv.table_name = @t ORDER BY rv.pk_json LIMIT @page OFFSET @off";
+    }
+
+    private static string Ident(string name) => "\"" + name.Replace("\"", "\"\"") + "\"";
+
+    private static string StringLiteral(string value) => "'" + value.Replace("'", "''") + "'";
+
+    // A single-quoted JSON path with a quoted key member so a column name with reserved characters
+    // still resolves: $."col".
+    private static string JsonPath(string column) => "'$.\"" + column.Replace("\"", "\\\"") + "\"'";
+}
+
+/// <summary>
+/// The inbound-snapshot consumer wired into <see cref="SyncSession"/>. Each batch merges through the
+/// shared <see cref="LwwApplier"/> (seq 0 — snapshot rows carry no oplog seq — so the applier's
+/// MAX-guarded watermark update is a no-op, and per-row LWW protects any newer local write). On
+/// complete, the remote watermark is set to the snapshot's as-of seq so deltas resume above it.
+/// </summary>
+internal sealed class SnapshotApplier : ISnapshotApplier
+{
+    private readonly LwwApplier _applier;
+    private readonly OplogStore _store;
+    private readonly ILogger _logger;
+    private long _asOfSeq;
+    private long _rows;
+    private long _maxHlc;
+
+    // db is part of the collaborator's canonical shape (mirrors the streamer) even though the merge
+    // reaches the database through the applier and store; kept for construction symmetry in the factory.
+    public SnapshotApplier(SqliteLocalDb db, LwwApplier applier, OplogStore store, ILogger logger)
+    {
+        _ = db;
+        _applier = applier;
+        _store = store;
+        _logger = logger;
+    }
+
+    public Task OnBeginAsync(SnapshotBegin begin, CancellationToken ct)
+    {
+        _asOfSeq = begin.AsOfSeq;
+        _rows = 0;
+        _maxHlc = 0;
+        _logger.LogInformation("Snapshot receive begun (as_of_seq {AsOfSeq}).", _asOfSeq);
+        return Task.CompletedTask;
+    }
+
+    public async Task OnBatchAsync(SnapshotBatch batch, CancellationToken ct)
+    {
+        if (batch.Rows.Count == 0)
+            return;
+
+        var entries = new List<OplogEntryRecord>(batch.Rows.Count);
+        foreach (var row in batch.Rows)
+        {
+            if (row.Hlc > _maxHlc) _maxHlc = row.Hlc;
+            entries.Add(new OplogEntryRecord(
+                0, batch.TableName, row.PkJson, row.HasRowJson ? row.RowJson : null, row.Hlc, row.NodeId, row.IsTombstone));
+        }
+        await _applier.ApplyBatchAsync(entries, ct);
+        _rows += batch.Rows.Count;
+    }
+
+    public async Task OnCompleteAsync(SnapshotComplete complete, CancellationToken ct)
+    {
+        // Preserve last_seen_hlc against a concurrently-applied delta that observed a higher HLC:
+        // SetLastAppliedRemoteSeqAsync overwrites the column, so pass the max of what we saw and
+        // what is already recorded.
+        var lastSeenHlc = Math.Max(_maxHlc, (await _store.GetPeerStateAsync(ct)).LastSeenHlc);
+        await _store.SetLastAppliedRemoteSeqAsync(complete.AsOfSeq, lastSeenHlc, ct);
+        _logger.LogInformation("Snapshot receive complete (as_of_seq {AsOfSeq}, {Rows} rows).", complete.AsOfSeq, _rows);
+    }
+}
@@ -39,6 +39,10 @@ internal sealed class SyncSession
    private string? _peerNodeId;
    private bool _driftWarned;

+    // Set true between an inbound SnapshotBegin and its SnapshotComplete. Only ever touched on the
+    // single receive-loop task, so no synchronization is needed.
+    private bool _snapshotReceiving;
+
    public SyncSession(
        SqliteLocalDb db, OplogStore store, LwwApplier applier, ReplicationOptions options, ILogger logger,
        Func<DateTimeOffset>? utcNow = null)
@@ -52,16 +56,18 @@ internal sealed class SyncSession
    }

    /// <summary>
-    /// Invoked when the peer's HandshakeAck requests a snapshot before deltas. Receives a gated
-    /// send delegate that enqueues on the bounded data lane (never the raw transport). Task 12 fills this in.
+    /// Invoked when THIS node owes the peer a snapshot (its own snapshot-required computation, not
+    /// the peer's) before deltas resume. Receives a gated send delegate that enqueues on the bounded
+    /// data lane (never the raw transport). Wired to <see cref="SnapshotStreamer.SendAsync"/>.
    /// </summary>
    public Func<Func<SyncMessage, CancellationToken, Task>, CancellationToken, Task>? SnapshotSender { get; set; }

    /// <summary>
-    /// Invoked on an inbound SnapshotBegin, with the triggering message and a gated data-lane send
-    /// delegate. Task 12 fills this in.
+    /// Consumes an inbound snapshot inline on the receive loop: SnapshotBegin opens the receive,
+    /// SnapshotBatch merges through LWW, SnapshotComplete advances the watermark. Wired to
+    /// <see cref="Internal.SnapshotApplier"/>.
    /// </summary>
-    public Func<SnapshotBegin, Func<SyncMessage, CancellationToken, Task>, CancellationToken, Task>? SnapshotReceiver { get; set; }
+    public ISnapshotApplier? SnapshotApplier { get; set; }

    /// <summary>Highest seq the peer has acknowledged applying this session.</summary>
    public long PeerAckedSeq => Interlocked.Read(ref _peerAckedSeq);
@@ -86,7 +92,7 @@ internal sealed class SyncSession
            await HandshakeAsync(duplex, control.Writer, GatedDataSend, linked.Token);

            var pump = PumpLoopAsync(data.Writer, linked.Token);
-            var receive = ReceiveLoopAsync(duplex, control.Writer, GatedDataSend, linked.Token);
+            var receive = ReceiveLoopAsync(duplex, control.Writer, linked.Token);
            var loops = new[] { pump, receive, writer };

            // The first COMPLETED task is the primary outcome. Cancel the siblings, observe every
@@ -142,7 +148,10 @@ internal sealed class SyncSession
            throw new InvalidOperationException(
                $"Replication protocol error: expected a HandshakeAck, got {peerAckMsg.MsgCase}.");

-        if (peerAckMsg.HandshakeAck.SnapshotRequired)
+        // We stream the snapshot when WE owe the peer one (computed from our pruned horizon /
+        // needs_snapshot), then it lands on the peer's receive loop. The peer's own flag drives ITS
+        // send direction symmetrically; both are safe to run because LWW merges either way.
+        if (snapshotRequired)
        {
            if (SnapshotSender is null)
                throw new NotSupportedException("snapshot required but no snapshot sender configured");
@@ -213,8 +222,7 @@ internal sealed class SyncSession
    }

    private async Task ReceiveLoopAsync(
-        SyncDuplex duplex, ChannelWriter<SyncMessage> controlLane,
-        Func<SyncMessage, CancellationToken, Task> dataSend, CancellationToken ct)
+        SyncDuplex duplex, ChannelWriter<SyncMessage> controlLane, CancellationToken ct)
    {
        while (await duplex.Inbox.WaitToReadAsync(ct))
        {
@@ -223,6 +231,7 @@ internal sealed class SyncSession
                switch (msg.MsgCase)
                {
                    case SyncMessage.MsgOneofCase.DeltaBatch:
+                        // Deltas apply normally even mid-snapshot: LWW makes the interleave safe.
                        await HandleDeltaBatchAsync(controlLane, msg.DeltaBatch, ct);
                        break;
                    case SyncMessage.MsgOneofCase.DeltaAck:
@@ -231,9 +240,23 @@ internal sealed class SyncSession
                        await _store.RecordPeerAckAsync(acked, ct);
                        break;
                    case SyncMessage.MsgOneofCase.SnapshotBegin:
-                        if (SnapshotReceiver is null)
-                            throw new NotSupportedException("snapshot received but no snapshot receiver configured");
-                        await SnapshotReceiver(msg.SnapshotBegin, dataSend, ct);
+                        if (SnapshotApplier is null)
+                            throw new NotSupportedException("snapshot received but no snapshot applier configured");
+                        _snapshotReceiving = true;
+                        await SnapshotApplier.OnBeginAsync(msg.SnapshotBegin, ct);
+                        break;
+                    case SyncMessage.MsgOneofCase.SnapshotBatch:
+                        if (!_snapshotReceiving)
+                            throw new InvalidOperationException(
+                                "Replication protocol error: SnapshotBatch with no active snapshot (missing SnapshotBegin).");
+                        await SnapshotApplier!.OnBatchAsync(msg.SnapshotBatch, ct);
+                        break;
+                    case SyncMessage.MsgOneofCase.SnapshotComplete:
+                        if (!_snapshotReceiving)
+                            throw new InvalidOperationException(
+                                "Replication protocol error: SnapshotComplete with no active snapshot (missing SnapshotBegin).");
+                        await SnapshotApplier!.OnCompleteAsync(msg.SnapshotComplete, ct);
+                        _snapshotReceiving = false;
                        break;
                    default:
                        throw new InvalidOperationException(
@@ -25,6 +25,11 @@ internal static class SyncSessionFactory
        var store = new OplogStore(db, options);
        var applier = new LwwApplier(db);
        var logger = loggerFactory.CreateLogger<SyncSession>();
-        return new SyncSession(db, store, applier, options, logger);
+        var session = new SyncSession(db, store, applier, options, logger);
+
+        var streamer = new SnapshotStreamer(db, store, options, loggerFactory.CreateLogger<SnapshotStreamer>());
+        session.SnapshotSender = streamer.SendAsync;
+        session.SnapshotApplier = new SnapshotApplier(db, applier, store, loggerFactory.CreateLogger<SnapshotApplier>());
+        return session;
    }
 }