scadaproj/ZB.MOM.WW.LocalDb/src/ZB.MOM.WW.LocalDb.Replication/Internal/SyncSession.cs

using System.Globalization;
using System.Threading.Channels;
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. All outbound traffic flows
/// through a single writer task fed by two lanes: an unbounded CONTROL lane (handshake messages,
/// acks) and a bounded DATA lane (delta/snapshot batches). Receiving therefore never blocks on the
/// ability to write — acks always fit — while the bounded data lane is the pump's backpressure.
/// Local writes never block on the peer; the first fault tears down all loops and propagates.
/// </summary>
internal sealed class SyncSession
{
    private const uint LibSchemaVersion = 1;

    // Data-lane depth bounds outbound batch buffering; control-lane volume is bounded by the
    // inbound batch rate (the peer's flow control), so its lane can safely be unbounded.
    private const int DataLaneCapacity = 16;

    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;

    // 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;
    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)
    {
        _db = db;
        _store = store;
        _applier = applier;
        _options = options;
        _logger = logger;
        _utcNow = utcNow ?? (() => DateTimeOffset.UtcNow);
    }

    /// <summary>
    /// 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); returns the as-of oplog seq the snapshot covers so the
    /// pump can skip already-covered tail deltas. Wired to <see cref="SnapshotStreamer.SendAsync"/>.
    /// </summary>
    public Func<Func<SyncMessage, CancellationToken, Task>, CancellationToken, Task<long>>? SnapshotSender { get; set; }

    /// <summary>
    /// 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 ISnapshotApplier? SnapshotApplier { get; set; }

    /// <summary>Highest seq the peer has acknowledged applying this session.</summary>
    public long PeerAckedSeq => Interlocked.Read(ref _peerAckedSeq);

    /// <summary>Highest oplog seq handed to the wire (or covered by a sent snapshot). Test/metric observation only.</summary>
    internal long SentThruSeq => Interlocked.Read(ref _sentThruSeq);

    /// <summary>The peer node's identity, once the handshake has completed. Volatile: written on the session task, read from test/metric threads.</summary>
    public string? PeerNodeId
    {
        get => Volatile.Read(ref _peerNodeId);
        private set => Volatile.Write(ref _peerNodeId, value);
    }

    public async Task RunAsync(SyncDuplex duplex, CancellationToken ct)
    {
        var control = Channel.CreateUnbounded<SyncMessage>(new UnboundedChannelOptions { SingleReader = true });
        var data = Channel.CreateBounded<SyncMessage>(new BoundedChannelOptions(DataLaneCapacity) { SingleReader = true });
        Task GatedDataSend(SyncMessage message, CancellationToken token) => data.Writer.WriteAsync(message, token).AsTask();

        using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct);
        var writer = WriterLoopAsync(duplex, control.Reader, data.Reader, linked.Token);
        try
        {
            await HandshakeAsync(duplex, control.Writer, GatedDataSend, linked.Token);

            var pump = PumpLoopAsync(data.Writer, 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
            // one of them under a swallow (so no secondary exception can replace the primary and
            // nothing is left running/unobserved), then surface the primary by awaiting it.
            var finished = await Task.WhenAny(loops);
            linked.Cancel();
            foreach (var loop in loops)
                if (!ReferenceEquals(loop, finished))
                    await DrainSecondaryAsync(loop);
            await finished;
        }
        catch
        {
            // Handshake-phase faults (and the primary rethrow above) land here: the writer may
            // still be running, so stop and observe it before the primary exception surfaces.
            linked.Cancel();
            await DrainSecondaryAsync(writer);
            throw;
        }
    }

    private async Task HandshakeAsync(
        SyncDuplex duplex, ChannelWriter<SyncMessage> controlLane,
        Func<SyncMessage, CancellationToken, Task> dataSend, CancellationToken ct)
    {
        var peerState = await _store.GetPeerStateAsync(ct);
        var localHandshake = BuildHandshake(peerState);

        // Both peers send before reading, so the transport must buffer at least two in-flight
        // messages (Handshake + HandshakeAck); gRPC's initial flow-control window covers that.
        await controlLane.WriteAsync(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 controlLane.WriteAsync(
            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}.");

        // 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.
        // Memory profile of a MUTUAL snapshot: while we send here (receive loop not yet running),
        // the peer's inbound frames — potentially its entire snapshot — buffer in the unbounded
        // duplex inbox until our send completes. Worst case one full peer snapshot held in memory;
        // acceptable for the small LocalDb datasets this library targets.
        if (snapshotRequired)
        {
            if (SnapshotSender is null)
                throw new NotSupportedException("snapshot required but no snapshot sender configured");
            var snapshotAsOfSeq = await SnapshotSender(dataSend, ct);
            // The snapshot already carries every row state up to as-of; skipping the pump past it
            // avoids re-sending tail deltas the peer would only discard via LWW.
            if (snapshotAsOfSeq > _sentThruSeq)
                _sentThruSeq = snapshotAsOfSeq;
        }
    }

    // The ONLY caller of duplex.Send. Control (handshake + acks) always preempts data so the peer's
    // inbound processing is never starved of its acks behind a backlog of DeltaBatches.
    private static async Task WriterLoopAsync(
        SyncDuplex duplex, ChannelReader<SyncMessage> control, ChannelReader<SyncMessage> data, CancellationToken ct)
    {
        Task<bool>? controlWait = null;
        Task<bool>? dataWait = null;
        var controlDone = false;
        var dataDone = false;
        while (true)
        {
            ct.ThrowIfCancellationRequested();
            if (control.TryRead(out var message) || data.TryRead(out message))
            {
                await duplex.Send(message, ct);
                continue;
            }
            // WaitToReadAsync => false means the lane was COMPLETED and is drained: stop waiting on
            // it (re-creating its waiter would instantly re-complete false — a 100% CPU spin) and
            // exit once both lanes are terminal. Nothing completes the lanes today, but Task 12 or
            // a future refactor must not silently detonate this loop.
            if (controlDone && dataDone)
                return;
            // At most one outstanding waiter per lane (SingleReader); a stale already-completed
            // waiter only costs one extra TryRead pass.
            if (!controlDone) controlWait ??= control.WaitToReadAsync(ct).AsTask();
            if (!dataDone) dataWait ??= data.WaitToReadAsync(ct).AsTask();
            await Task.WhenAny(controlWait ?? dataWait!, dataWait ?? controlWait!);
            if (controlWait is { IsCompleted: true })
            {
                controlDone = !await controlWait;
                controlWait = null;
            }
            if (dataWait is { IsCompleted: true })
            {
                dataDone = !await dataWait;
                dataWait = null;
            }
        }
    }

    private async Task PumpLoopAsync(ChannelWriter<SyncMessage> dataLane, CancellationToken ct)
    {
        while (true)
        {
            ct.ThrowIfCancellationRequested();
            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));
                // Blocking here on a full data lane IS the pump's backpressure.
                await dataLane.WriteAsync(new SyncMessage { DeltaBatch = delta }, ct);
                _sentThruSeq = batch[^1].Seq;
                continue;
            }

            await Task.Delay(_options.FlushInterval, ct);
        }
    }

    private async Task ReceiveLoopAsync(
        SyncDuplex duplex, ChannelWriter<SyncMessage> controlLane, CancellationToken ct)
    {
        while (await duplex.Inbox.WaitToReadAsync(ct))
        {
            while (duplex.Inbox.TryRead(out var msg))
            {
                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:
                        var acked = msg.DeltaAck.AppliedThruSeq;
                        InterlockedMax(ref _peerAckedSeq, acked);
                        await _store.RecordPeerAckAsync(acked, ct);
                        break;
                    case SyncMessage.MsgOneofCase.SnapshotBegin:
                        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(
                            $"Replication protocol error: unexpected {msg.MsgCase} after the handshake.");
                }
            }
        }
    }

    private async Task HandleDeltaBatchAsync(ChannelWriter<SyncMessage> controlLane, DeltaBatch batch, CancellationToken ct)
    {
        if (batch.Entries.Count == 0)
            return;

        var now = _utcNow();
        var nowMs = now.ToUnixTimeMilliseconds();
        var driftLimitMs = nowMs + (long)_options.MaxHlcDriftAhead.TotalMilliseconds;
        var toApply = new List<OplogEntryRecord>(batch.Entries.Count);
        var maxSeq = 0L;

        foreach (var proto in batch.Entries)
        {
            if (proto.Seq > maxSeq) maxSeq = proto.Seq;
            var entry = FromProto(proto);
            var physicalMs = HybridLogicalClock.PhysicalMs(entry.Hlc);
            if (physicalMs > driftLimitMs)
            {
                if (!_driftWarned)
                {
                    // Rate-limited to one warning per session: drift is usually a sustained clock
                    // condition and would otherwise flood the log at the delta rate.
                    _driftWarned = true;
                    _logger.LogWarning(
                        "Inbound oplog entry HLC physical time {EntryMs} ms is ahead of local UTC {NowMs} ms by more than MaxHlcDriftAhead ({Drift}). Further drift this session is not re-logged.",
                        physicalMs, nowMs, _options.MaxHlcDriftAhead);
                }
                // Dead-letter is the fail-closed disposition: faulting the session instead would
                // livelock on reconnect (the peer resends the same drifted entries forever).
                if (_options.FailClosedOnDrift)
                {
                    await DeadLetterDriftedAsync(entry, physicalMs, nowMs, now, ct);
                    continue;
                }
            }
            toApply.Add(entry);
        }

        var result = await _applier.ApplyBatchAsync(toApply, ct);
        var ackThru = Math.Max(result.AppliedThruSeq, maxSeq);
        if (ackThru > result.AppliedThruSeq)
            // Dead-lettered entries carried the batch's top seqs; advance the durable remote
            // watermark over them (MAX-guarded like the applier's) so a reconnect handshake does
            // not make the peer resend what this node has already discarded.
            await _db.ExecuteAsync(
                "UPDATE __localdb_peer_state SET last_applied_remote_seq = MAX(last_applied_remote_seq, @seq) WHERE id = 1",
                new { seq = ackThru }, ct);
        await controlLane.WriteAsync(new SyncMessage { DeltaAck = new DeltaAck { AppliedThruSeq = ackThru } }, ct);
    }

    // Not transactional with the watermark update: a crash between them redelivers the batch and
    // dead-letters the same entry again (at-least-once — a benign duplicate dead-letter row).
    private Task DeadLetterDriftedAsync(
        OplogEntryRecord entry, long physicalMs, long nowMs, DateTimeOffset now, CancellationToken ct) =>
        _db.ExecuteAsync(
            "INSERT INTO __localdb_dead_letter (received_utc, table_name, pk_json, row_json, hlc, node_id, error) " +
            "VALUES (@utc, @t, @pk, @row, @hlc, @node, @err)",
            new
            {
                utc = now.UtcDateTime.ToString(OplogStore.TombstoneUtcFormat, CultureInfo.InvariantCulture),
                t = entry.TableName,
                pk = entry.PkJson,
                row = entry.RowJson,
                hlc = entry.Hlc,
                node = entry.NodeId,
                err = $"HLC drift ahead: entry physical time {physicalMs} ms (Unix UTC) exceeds local now {nowMs} ms + MaxHlcDriftAhead {_options.MaxHlcDriftAhead} (fail-closed).",
            }, ct);

    private async Task DrainSecondaryAsync(Task task)
    {
        try
        {
            await task;
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Secondary sync-session task ended during teardown.");
        }
    }

    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 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);
    }
}