feat(localdb): SyncSession protocol state machine (handshake fail-closed, bidirectional delta pump)

This commit is contained in:
Joseph Doherty
2026-07-17 22:26:03 -04:00
parent b223507070
commit 925623b878
3 changed files with 750 additions and 0 deletions
@@ -0,0 +1,18 @@
using System.Threading.Channels;
using ZB.MOM.WW.LocalDb.Contracts;
namespace ZB.MOM.WW.LocalDb.Replication.Internal;
/// <summary>
/// The transport seam a <see cref="SyncSession"/> runs over: an outbound send delegate plus an
/// inbound message reader. Task 11 adapts a gRPC duplex stream to this; tests wire two in-memory
/// <see cref="Channel{T}"/> pairs crosswise. <see cref="Send"/> may be called concurrently by the
/// session's pump and receive loops — the session serializes those calls, so an adapter's writer
/// only ever sees one send at a time.
/// </summary>
internal sealed class SyncDuplex
{
public required Func<SyncMessage, CancellationToken, Task> Send { get; init; }
public required ChannelReader<SyncMessage> Inbox { get; init; }
}
@@ -0,0 +1,286 @@
using Microsoft.Extensions.Logging;
using ZB.MOM.WW.LocalDb.Contracts;
using ZB.MOM.WW.LocalDb.Hlc;
using ZB.MOM.WW.LocalDb.Internal;
namespace ZB.MOM.WW.LocalDb.Replication.Internal;
/// <summary>
/// The symmetric 2-node sync protocol as one state machine: both the gRPC service and the gRPC
/// client run this identical logic over a <see cref="SyncDuplex"/>. It performs the fail-closed
/// handshake, then runs a bidirectional steady state — an outbound change pump and an inbound
/// apply/ack loop — until the token is cancelled or the stream ends. Local writes never block on
/// the peer; every failure is a typed exception that tears down both loops and propagates.
/// </summary>
internal sealed class SyncSession
{
private const uint LibSchemaVersion = 1;
private readonly SqliteLocalDb _db;
private readonly OplogStore _store;
private readonly LwwApplier _applier;
private readonly ReplicationOptions _options;
private readonly ILogger _logger;
private readonly Func<DateTimeOffset> _utcNow;
// gRPC stream writers are not safe for concurrent writes; the pump and receive loops both send,
// so every outbound message goes through this gate to serialize them into one writer.
private readonly SemaphoreSlim _sendGate = new(1, 1);
// Highest oplog seq handed to the wire this session; the pump reads strictly above it so a batch
// is never re-sent before its ack (the reliable ordered stream delivers exactly once in-session).
private long _sentThruSeq;
private long _peerAckedSeq;
public SyncSession(
SqliteLocalDb db, OplogStore store, LwwApplier applier, ReplicationOptions options, ILogger logger,
Func<DateTimeOffset>? utcNow = null)
{
_db = db;
_store = store;
_applier = applier;
_options = options;
_logger = logger;
_utcNow = utcNow ?? (() => DateTimeOffset.UtcNow);
}
/// <summary>Invoked when the peer's HandshakeAck requests a snapshot before deltas. Task 12 fills this in.</summary>
public Func<SyncDuplex, CancellationToken, Task>? SnapshotSender { get; set; }
/// <summary>Invoked on an inbound SnapshotBegin. Task 12 fills this in.</summary>
public Func<SnapshotBegin, SyncDuplex, CancellationToken, Task>? SnapshotReceiver { get; set; }
/// <summary>Highest seq the peer has acknowledged applying this session.</summary>
public long PeerAckedSeq => Interlocked.Read(ref _peerAckedSeq);
/// <summary>The peer node's identity, once the handshake has completed.</summary>
public string? PeerNodeId { get; private set; }
public async Task RunAsync(SyncDuplex duplex, CancellationToken ct)
{
var peerState = await _store.GetPeerStateAsync(ct);
var localHandshake = BuildHandshake(peerState);
await SendAsync(duplex, new SyncMessage { Handshake = localHandshake }, ct);
var firstInbound = await ReceiveOneAsync(duplex, ct);
if (firstInbound.MsgCase != SyncMessage.MsgOneofCase.Handshake)
throw new InvalidOperationException(
$"Replication protocol error: expected a Handshake as the first message, got {firstInbound.MsgCase}.");
var peerHandshake = firstInbound.Handshake;
ValidateHandshake(localHandshake, peerHandshake);
await _store.SetPeerNodeIdAsync(peerHandshake.NodeId, ct);
PeerNodeId = peerHandshake.NodeId;
_sentThruSeq = peerHandshake.LastAppliedRemoteSeq;
Interlocked.Exchange(ref _peerAckedSeq, peerHandshake.LastAppliedRemoteSeq);
var snapshotRequired = await ComputeSnapshotRequiredAsync(peerHandshake, peerState, ct);
await SendAsync(duplex, new SyncMessage { HandshakeAck = new HandshakeAck { NodeId = _db.NodeId, SnapshotRequired = snapshotRequired } }, ct);
var peerAckMsg = await ReceiveOneAsync(duplex, ct);
if (peerAckMsg.MsgCase != SyncMessage.MsgOneofCase.HandshakeAck)
throw new InvalidOperationException(
$"Replication protocol error: expected a HandshakeAck, got {peerAckMsg.MsgCase}.");
if (peerAckMsg.HandshakeAck.SnapshotRequired)
{
if (SnapshotSender is null)
throw new NotSupportedException("snapshot required but no snapshot sender configured");
await SnapshotSender(duplex, ct);
}
await RunSteadyStateAsync(duplex, ct);
}
private async Task RunSteadyStateAsync(SyncDuplex duplex, CancellationToken ct)
{
using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct);
var pump = PumpLoopAsync(duplex, linked.Token);
var receive = ReceiveLoopAsync(duplex, linked.Token);
// Whichever loop completes first (fault or stream-end) drives teardown: cancel the sibling,
// drain it (swallowing only the cancellation we caused), then surface the primary outcome.
var finished = await Task.WhenAny(pump, receive);
linked.Cancel();
var other = finished == pump ? receive : pump;
try { await other; }
catch (OperationCanceledException) { }
await finished;
}
private async Task PumpLoopAsync(SyncDuplex duplex, CancellationToken ct)
{
while (!ct.IsCancellationRequested)
{
var batch = await _store.ReadBatchAboveAsync(_sentThruSeq, _options.MaxBatchSize, ct);
if (batch.Count > 0)
{
var delta = new DeltaBatch();
foreach (var entry in batch)
delta.Entries.Add(ToProto(entry));
await SendAsync(duplex, new SyncMessage { DeltaBatch = delta }, ct);
_sentThruSeq = batch[^1].Seq;
continue;
}
await Task.Delay(_options.FlushInterval, ct);
}
}
private async Task ReceiveLoopAsync(SyncDuplex duplex, CancellationToken ct)
{
while (await duplex.Inbox.WaitToReadAsync(ct))
{
while (duplex.Inbox.TryRead(out var msg))
{
switch (msg.MsgCase)
{
case SyncMessage.MsgOneofCase.DeltaBatch:
await HandleDeltaBatchAsync(duplex, msg.DeltaBatch, ct);
break;
case SyncMessage.MsgOneofCase.DeltaAck:
var acked = msg.DeltaAck.AppliedThruSeq;
InterlockedMax(ref _peerAckedSeq, acked);
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, duplex, ct);
break;
default:
throw new InvalidOperationException(
$"Replication protocol error: unexpected {msg.MsgCase} after the handshake.");
}
}
}
}
private async Task HandleDeltaBatchAsync(SyncDuplex duplex, DeltaBatch batch, CancellationToken ct)
{
var entries = new List<OplogEntryRecord>(batch.Entries.Count);
var maxHlc = 0L;
foreach (var e in batch.Entries)
{
entries.Add(FromProto(e));
if (e.Hlc > maxHlc) maxHlc = e.Hlc;
}
if (entries.Count > 0)
{
var maxHlcMs = HybridLogicalClock.PhysicalMs(maxHlc);
var driftLimitMs = _utcNow().Add(_options.MaxHlcDriftAhead).ToUnixTimeMilliseconds();
if (maxHlcMs > driftLimitMs)
{
_logger.LogWarning(
"Inbound delta HLC physical time {HlcMs} ms exceeds local clock + MaxHlcDriftAhead ({LimitMs} ms).",
maxHlcMs, driftLimitMs);
if (_options.FailClosedOnDrift)
throw new InvalidOperationException(
$"Inbound HLC physical time {maxHlcMs} ms exceeds MaxHlcDriftAhead limit {driftLimitMs} ms; rejecting batch (fail-closed).");
}
}
var result = await _applier.ApplyBatchAsync(entries, ct);
await SendAsync(duplex, new SyncMessage { DeltaAck = new DeltaAck { AppliedThruSeq = result.AppliedThruSeq } }, ct);
}
private Handshake BuildHandshake(PeerState peerState)
{
var handshake = new Handshake
{
NodeId = _db.NodeId,
LibSchemaVersion = LibSchemaVersion,
LastAppliedRemoteSeq = peerState.LastAppliedRemoteSeq,
};
foreach (var table in _db.ReplicatedTables.Values.OrderBy(t => t.Name, StringComparer.Ordinal))
handshake.Tables.Add(new TableDigest { TableName = table.Name, Digest = table.Digest });
return handshake;
}
private static void ValidateHandshake(Handshake local, Handshake peer)
{
var versionMismatch = local.LibSchemaVersion != peer.LibSchemaVersion;
var localByName = local.Tables.ToDictionary(t => t.TableName, t => t.Digest, StringComparer.Ordinal);
var peerByName = peer.Tables.ToDictionary(t => t.TableName, t => t.Digest, StringComparer.Ordinal);
var mismatched = new SortedSet<string>(StringComparer.Ordinal);
foreach (var name in localByName.Keys.Union(peerByName.Keys, StringComparer.Ordinal))
if (!localByName.TryGetValue(name, out var ld) || !peerByName.TryGetValue(name, out var pd) || ld != pd)
mismatched.Add(name);
if (!versionMismatch && mismatched.Count == 0)
return;
var parts = new List<string>();
if (versionMismatch)
parts.Add($"lib_schema_version {local.LibSchemaVersion} != peer {peer.LibSchemaVersion}");
if (mismatched.Count > 0)
parts.Add($"table schema mismatch on: {string.Join(", ", mismatched)}");
throw new LocalDbSchemaMismatchException(
"Replication handshake rejected (fail-closed): " + string.Join("; ", parts) + ".");
}
private async Task<bool> ComputeSnapshotRequiredAsync(Handshake peerHandshake, PeerState peerState, CancellationToken ct)
{
if (peerState.NeedsSnapshot)
return true;
var minRows = await _db.QueryAsync(
"SELECT COALESCE(MIN(seq), 0) FROM __localdb_oplog", static r => r.GetInt64(0), null, ct);
var minSeq = minRows[0];
// minSeq 0 => oplog empty. A gap exists when the peer's applied watermark falls below the
// oldest seq we can still stream (pruning discarded everything at or below its horizon).
return minSeq > 0 && peerHandshake.LastAppliedRemoteSeq < minSeq - 1;
}
private async Task SendAsync(SyncDuplex duplex, SyncMessage message, CancellationToken ct)
{
await _sendGate.WaitAsync(ct);
try { await duplex.Send(message, ct); }
finally { _sendGate.Release(); }
}
private static async Task<SyncMessage> ReceiveOneAsync(SyncDuplex duplex, CancellationToken ct)
{
if (await duplex.Inbox.WaitToReadAsync(ct) && duplex.Inbox.TryRead(out var message))
return message;
throw new InvalidOperationException(
"Replication protocol error: the peer closed the stream before completing the handshake.");
}
private static OplogEntry ToProto(OplogEntryRecord entry)
{
var proto = new OplogEntry
{
Seq = entry.Seq,
TableName = entry.TableName,
PkJson = entry.PkJson,
Hlc = entry.Hlc,
NodeId = entry.NodeId,
IsTombstone = entry.IsTombstone,
};
if (entry.RowJson is not null)
proto.RowJson = entry.RowJson;
return proto;
}
private static OplogEntryRecord FromProto(OplogEntry entry) =>
new(entry.Seq, entry.TableName, entry.PkJson, entry.HasRowJson ? entry.RowJson : null,
entry.Hlc, entry.NodeId, entry.IsTombstone);
private static void InterlockedMax(ref long target, long value)
{
long current;
do
{
current = Interlocked.Read(ref target);
if (value <= current)
return;
}
while (Interlocked.CompareExchange(ref target, value, current) != current);
}
}