scadaproj/ZB.MOM.WW.LocalDb/tests/ZB.MOM.WW.LocalDb.Tests/BatchBudgetTests.cs

using System.Threading.Channels;
using Microsoft.Data.Sqlite;
using Microsoft.Extensions.Logging.Abstractions;
using ZB.MOM.WW.LocalDb.Contracts;
using ZB.MOM.WW.LocalDb.Internal;
using ZB.MOM.WW.LocalDb.Replication;
using ZB.MOM.WW.LocalDb.Replication.Internal;

namespace ZB.MOM.WW.LocalDb.Tests;

/// <summary>
/// Outbound batching by summed serialized BYTES (with the row count kept as a secondary cap), and
/// the inbound coalescing that keeps one transaction + one ack per burst rather than per message.
/// </summary>
public sealed class BatchBudgetTests : IDisposable
{
    private static readonly TimeSpan RunTimeout = TimeSpan.FromSeconds(15);

    private readonly List<string> _paths = [];
    private readonly List<IDisposable> _disposables = [];

    public void Dispose()
    {
        foreach (var d in _disposables)
            d.Dispose();
        SqliteConnection.ClearAllPools();
        foreach (var p in _paths)
        {
            if (File.Exists(p)) File.Delete(p);
            if (File.Exists(p + "-wal")) File.Delete(p + "-wal");
            if (File.Exists(p + "-shm")) File.Delete(p + "-shm");
        }
    }

    private async Task<SqliteLocalDb> NewOrdersDbAsync()
    {
        var path = Path.Combine(Path.GetTempPath(), Guid.NewGuid() + ".db");
        _paths.Add(path);
        var db = new SqliteLocalDb(new LocalDbOptions { Path = path });
        _disposables.Add(db);
        await db.ExecuteAsync("CREATE TABLE orders (id INTEGER PRIMARY KEY, sku TEXT, qty INTEGER)");
        db.RegisterReplicated("orders");
        return db;
    }

    private static (SyncDuplex Duplex, ChannelWriter<SyncMessage> ToSession, ChannelReader<SyncMessage> FromSession) ScriptedPeer()
    {
        var toSession = Channel.CreateUnbounded<SyncMessage>();
        var fromSession = Channel.CreateUnbounded<SyncMessage>();
        var duplex = new SyncDuplex
        {
            Send = (m, ct) => fromSession.Writer.WriteAsync(m, ct).AsTask(),
            Inbox = toSession.Reader,
        };
        return (duplex, toSession.Writer, fromSession.Reader);
    }

    private static Handshake MatchingHandshake(SqliteLocalDb sessionDb, string peerNodeId)
    {
        var h = new Handshake { NodeId = peerNodeId, LibSchemaVersion = 1, LastAppliedRemoteSeq = 0 };
        foreach (var t in sessionDb.ReplicatedTables.Values.OrderBy(t => t.Name, StringComparer.Ordinal))
            h.Tables.Add(new TableDigest { TableName = t.Name, Digest = t.Digest });
        return h;
    }

    private static async Task<SyncMessage> NextAsync(ChannelReader<SyncMessage> reader, CancellationToken ct)
    {
        await reader.WaitToReadAsync(ct);
        reader.TryRead(out var msg);
        return msg!;
    }

    private static async Task SwallowAsync(Task task)
    {
        try { await task; }
        catch (OperationCanceledException) { }
    }

    // Queues a complete peer handshake AHEAD of the session's own, so anything queued after it is
    // guaranteed to be sitting in the inbox by the time the receive loop starts. That is what makes
    // the coalescing assertions deterministic rather than a race with the reader.
    private static async Task PreloadHandshakeAsync(
        SqliteLocalDb db, ChannelWriter<SyncMessage> toSession, CancellationToken ct)
    {
        await toSession.WriteAsync(new SyncMessage { Handshake = MatchingHandshake(db, "peer") }, ct);
        await toSession.WriteAsync(
            new SyncMessage { HandshakeAck = new HandshakeAck { NodeId = "peer", SnapshotRequired = false } }, ct);
    }

    private static long BatchBytes(DeltaBatch batch) =>
        batch.Entries.Sum(e => 1L + Google.Protobuf.CodedOutputStream.ComputeMessageSize(e));

    private static long BatchBytes(SnapshotBatch batch) =>
        batch.Rows.Sum(r => 1L + Google.Protobuf.CodedOutputStream.ComputeMessageSize(r));

    private static SyncSession NewSession(SqliteLocalDb db, ReplicationOptions options)
    {
        var store = new OplogStore(db, options);
        var applier = new LwwApplier(db);
        return new SyncSession(db, store, applier, options, NullLogger.Instance);
    }

    private static OplogEntry ProtoEntry(long seq, long id, long hlc) =>
        new()
        {
            Seq = seq,
            TableName = "orders",
            PkJson = $"{{\"id\":{id}}}",
            RowJson = $"{{\"id\":{id},\"sku\":\"S\",\"qty\":1}}",
            Hlc = hlc,
            NodeId = "peer",
            IsTombstone = false,
        };

    // ---- Outbound: delta pump --------------------------------------------------------------

    [Fact]
    public async Task Pump_SplitsOneReadIntoSeveralMessages_EachWithinTheByteBudget()
    {
        const int budget = 4_000;
        var db = await NewOrdersDbAsync();
        var wide = new string('x', 900);
        for (var i = 1; i <= 20; i++)
            await db.ExecuteAsync(
                "INSERT INTO orders (id, sku, qty) VALUES (@id, @sku, @id)", new { id = i, sku = wide });

        // Row count alone would put all 20 wide rows (~18 KB) in ONE message; the byte budget is
        // what actually bounds the message the transport has to carry.
        var options = new ReplicationOptions
        {
            FlushInterval = TimeSpan.FromMilliseconds(20), MaxBatchSize = 500, MaxBatchBytes = budget,
        };
        var session = NewSession(db, options);
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);
        await PreloadHandshakeAsync(db, toSession, cts.Token);

        var seen = new List<long>();
        while (seen.Count < 20)
        {
            var msg = await NextAsync(fromSession, cts.Token);
            if (msg.MsgCase != SyncMessage.MsgOneofCase.DeltaBatch)
                continue;
            Assert.True(BatchBytes(msg.DeltaBatch) <= budget,
                $"a {BatchBytes(msg.DeltaBatch)}-byte message exceeded the {budget}-byte budget");
            Assert.True(msg.DeltaBatch.Entries.Count > 1, "the budget should still fit several of these rows");
            seen.AddRange(msg.DeltaBatch.Entries.Select(e => e.Seq));
        }

        // Split, not dropped or duplicated: every captured seq travelled exactly once, in order.
        Assert.Equal(20, seen.Count);
        Assert.Equal(seen.OrderBy(x => x).ToList(), seen);
        Assert.Equal(20, seen.Distinct().Count());

        cts.Cancel();
        await SwallowAsync(run);
    }

    [Fact]
    public async Task Pump_RowLargerThanTheBudget_IsSentAloneRatherThanStallingTheStream()
    {
        var db = await NewOrdersDbAsync();
        await db.ExecuteAsync(
            "INSERT INTO orders (id, sku, qty) VALUES (1, @sku, 1)", new { sku = new string('x', 5_000) });
        await db.ExecuteAsync("INSERT INTO orders (id, sku, qty) VALUES (2, 'small', 2)");

        var options = new ReplicationOptions
        {
            FlushInterval = TimeSpan.FromMilliseconds(20), MaxBatchSize = 500, MaxBatchBytes = 100,
        };
        var session = NewSession(db, options);
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);
        await PreloadHandshakeAsync(db, toSession, cts.Token);

        var batches = new List<DeltaBatch>();
        while (batches.Sum(b => b.Entries.Count) < 2)
        {
            var msg = await NextAsync(fromSession, cts.Token);
            if (msg.MsgCase == SyncMessage.MsgOneofCase.DeltaBatch)
                batches.Add(msg.DeltaBatch);
        }

        // The oversized row goes out on its own — refusing it would wedge the stream on one row
        // forever, and everything behind it with it.
        Assert.Equal(2, batches.Count);
        Assert.Single(batches[0].Entries);
        Assert.Equal(1, batches[0].Entries[0].Seq);
        Assert.True(BatchBytes(batches[0]) > options.MaxBatchBytes);
        Assert.Single(batches[1].Entries);
        Assert.Equal(2, batches[1].Entries[0].Seq);

        cts.Cancel();
        await SwallowAsync(run);
    }

    [Fact]
    public async Task Pump_RowCountRemainsASecondaryCap()
    {
        var db = await NewOrdersDbAsync();
        for (var i = 1; i <= 12; i++)
            await db.ExecuteAsync("INSERT INTO orders (id, sku, qty) VALUES (@id, 'S', @id)", new { id = i });

        // Tiny rows, enormous byte budget: only the row cap can be doing the splitting here.
        var options = new ReplicationOptions
        {
            FlushInterval = TimeSpan.FromMilliseconds(20), MaxBatchSize = 5, MaxBatchBytes = 8 * 1024 * 1024,
        };
        var session = NewSession(db, options);
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);
        await PreloadHandshakeAsync(db, toSession, cts.Token);

        var seen = 0;
        while (seen < 12)
        {
            var msg = await NextAsync(fromSession, cts.Token);
            if (msg.MsgCase != SyncMessage.MsgOneofCase.DeltaBatch)
                continue;
            Assert.True(msg.DeltaBatch.Entries.Count <= 5);
            seen += msg.DeltaBatch.Entries.Count;
        }

        cts.Cancel();
        await SwallowAsync(run);
    }

    // ---- Outbound: snapshot --------------------------------------------------------------

    [Fact]
    public async Task Snapshot_SplitsPagesIntoMessagesWithinTheByteBudget()
    {
        const int budget = 4_000;
        var db = await NewOrdersDbAsync();
        var wide = new string('x', 900);
        for (var i = 1; i <= 20; i++)
            await db.ExecuteAsync(
                "INSERT INTO orders (id, sku, qty) VALUES (@id, @sku, @id)", new { id = i, sku = wide });

        // A snapshot is the easiest way to blow a message limit: it ships whole rows, all of them,
        // back to back — so its page size has to be bounded in bytes too, not just in rows.
        var options = new ReplicationOptions { MaxBatchSize = 500, MaxBatchBytes = budget };
        var store = new OplogStore(db, options);
        var streamer = new SnapshotStreamer(db, store, options, NullLogger.Instance);

        var sent = new List<SyncMessage>();
        await streamer.SendAsync(SnapshotScope.Full, (m, _) => { sent.Add(m); return Task.CompletedTask; }, default);

        var batches = sent.Where(m => m.MsgCase == SyncMessage.MsgOneofCase.SnapshotBatch)
            .Select(m => m.SnapshotBatch).ToList();
        Assert.True(batches.Count > 1, "20 wide rows must not have travelled as one message");
        Assert.All(batches, b => Assert.True(BatchBytes(b) <= budget,
            $"a {BatchBytes(b)}-byte snapshot message exceeded the {budget}-byte budget"));

        // Bracketed and complete: every row, exactly once.
        Assert.Equal(SyncMessage.MsgOneofCase.SnapshotBegin, sent[0].MsgCase);
        Assert.Equal(SyncMessage.MsgOneofCase.SnapshotComplete, sent[^1].MsgCase);
        var pks = batches.SelectMany(b => b.Rows).Select(r => r.PkJson).ToList();
        Assert.Equal(20, pks.Count);
        Assert.Equal(20, pks.Distinct().Count());
    }

    [Fact]
    public async Task Snapshot_RowLargerThanTheBudget_IsSentAlone()
    {
        var db = await NewOrdersDbAsync();
        await db.ExecuteAsync(
            "INSERT INTO orders (id, sku, qty) VALUES (1, @sku, 1)", new { sku = new string('x', 5_000) });
        await db.ExecuteAsync("INSERT INTO orders (id, sku, qty) VALUES (2, 'small', 2)");

        var options = new ReplicationOptions { MaxBatchSize = 500, MaxBatchBytes = 100 };
        var store = new OplogStore(db, options);
        var streamer = new SnapshotStreamer(db, store, options, NullLogger.Instance);

        var sent = new List<SyncMessage>();
        await streamer.SendAsync(SnapshotScope.Full, (m, _) => { sent.Add(m); return Task.CompletedTask; }, default);

        var batches = sent.Where(m => m.MsgCase == SyncMessage.MsgOneofCase.SnapshotBatch)
            .Select(m => m.SnapshotBatch).ToList();
        Assert.Equal(2, batches.Count);
        Assert.All(batches, b => Assert.Single(b.Rows));
        Assert.Equal("{\"id\":1}", batches[0].Rows[0].PkJson);
        Assert.Equal("{\"id\":2}", batches[1].Rows[0].PkJson);
    }

    // ---- Inbound: coalescing -------------------------------------------------------------

    [Fact]
    public async Task Receive_CoalescesQueuedDeltaBatches_IntoOneAck()
    {
        var db = await NewOrdersDbAsync();
        var options = new ReplicationOptions { FlushInterval = TimeSpan.FromSeconds(30) };
        var session = NewSession(db, options);
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);

        // Queued ahead of the receive loop, so all three are demonstrably in the inbox together.
        await PreloadHandshakeAsync(db, toSession, cts.Token);
        for (var i = 1; i <= 3; i++)
        {
            var batch = new DeltaBatch();
            batch.Entries.Add(ProtoEntry(i, i, 1_000_000 + i));
            await toSession.WriteAsync(new SyncMessage { DeltaBatch = batch }, cts.Token);
        }

        SyncMessage ack;
        do { ack = await NextAsync(fromSession, cts.Token); }
        while (ack.MsgCase != SyncMessage.MsgOneofCase.DeltaAck);

        // The discriminator: uncoalesced, the FIRST ack could only cover seq 1.
        Assert.Equal(3, ack.DeltaAck.AppliedThruSeq);

        var rows = await db.QueryAsync("SELECT COUNT(*) FROM orders", r => r.GetInt64(0));
        Assert.Equal(3, rows[0]);

        cts.Cancel();
        await SwallowAsync(run);
    }

    [Fact]
    public async Task Receive_DeltaCoalescing_DoesNotSwallowTheMessageThatEndsTheRun()
    {
        // The risk the one-message lookahead exists to handle: coalescing reads ahead, and whatever
        // it reads that is NOT a delta must still be processed, not dropped on the floor.
        var db = await NewOrdersDbAsync();
        var options = new ReplicationOptions { FlushInterval = TimeSpan.FromSeconds(30) };
        var session = NewSession(db, options);
        var captured = new CapturingSnapshotApplier();
        session.SnapshotApplier = captured;
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);

        await PreloadHandshakeAsync(db, toSession, cts.Token);
        var batch = new DeltaBatch();
        batch.Entries.Add(ProtoEntry(1, 1, 1_000_001));
        await toSession.WriteAsync(new SyncMessage { DeltaBatch = batch }, cts.Token);
        await toSession.WriteAsync(
            new SyncMessage { SnapshotBegin = new SnapshotBegin { AsOfSeq = 77 } }, cts.Token);

        SyncMessage ack;
        do { ack = await NextAsync(fromSession, cts.Token); }
        while (ack.MsgCase != SyncMessage.MsgOneofCase.DeltaAck);
        Assert.Equal(1, ack.DeltaAck.AppliedThruSeq);

        // The SnapshotBegin that terminated the coalescing run still reached its handler.
        await WaitForAsync(() => captured.BeginAsOfSeq == 77, RunTimeout);

        cts.Cancel();
        await SwallowAsync(run);
    }

    [Fact]
    public async Task Receive_CoalescesQueuedAcks_ToTheHighestWatermark_WithoutLosingTheNextMessage()
    {
        var db = await NewOrdersDbAsync();
        for (var i = 1; i <= 9; i++)
            await db.ExecuteAsync("INSERT INTO orders (id, sku, qty) VALUES (@id, 'S', @id)", new { id = i });

        var options = new ReplicationOptions { FlushInterval = TimeSpan.FromSeconds(30) };
        var session = NewSession(db, options);
        var captured = new CapturingSnapshotApplier();
        session.SnapshotApplier = captured;
        var (duplex, toSession, fromSession) = ScriptedPeer();

        using var cts = new CancellationTokenSource(RunTimeout);
        var run = session.RunAsync(duplex, cts.Token);

        await PreloadHandshakeAsync(db, toSession, cts.Token);
        // Out of order on purpose: collapsing acks must take the MAX, never the last.
        foreach (var seq in new long[] { 4, 9, 2 })
            await toSession.WriteAsync(
                new SyncMessage { DeltaAck = new DeltaAck { AppliedThruSeq = seq } }, cts.Token);
        await toSession.WriteAsync(
            new SyncMessage { SnapshotBegin = new SnapshotBegin { AsOfSeq = 55 } }, cts.Token);

        await WaitForAsync(async () =>
        {
            var acked = await db.QueryAsync(
                "SELECT last_acked_seq FROM __localdb_peer_state WHERE id = 1", r => r.GetInt64(0));
            return acked[0] == 9;
        }, RunTimeout);

        Assert.Equal(9, session.PeerAckedSeq);
        await WaitForAsync(() => captured.BeginAsOfSeq == 55, RunTimeout);

        cts.Cancel();
        await SwallowAsync(run);
    }

    private sealed class CapturingSnapshotApplier : ISnapshotApplier
    {
        private long _beginAsOfSeq = -1;

        public long BeginAsOfSeq => Interlocked.Read(ref _beginAsOfSeq);

        public Task OnBeginAsync(SnapshotBegin begin, CancellationToken ct)
        {
            Interlocked.Exchange(ref _beginAsOfSeq, begin.AsOfSeq);
            return Task.CompletedTask;
        }

        public Task OnBatchAsync(SnapshotBatch batch, CancellationToken ct) => Task.CompletedTask;

        public Task OnCompleteAsync(SnapshotComplete complete, CancellationToken ct) => Task.CompletedTask;
    }

    private static Task WaitForAsync(Func<bool> predicate, TimeSpan timeout) =>
        WaitForAsync(() => Task.FromResult(predicate()), timeout);

    private static async Task WaitForAsync(Func<Task<bool>> predicate, TimeSpan timeout)
    {
        var deadline = DateTime.UtcNow + timeout;
        while (DateTime.UtcNow < deadline)
        {
            if (await predicate()) return;
            await Task.Delay(20);
        }
        throw new TimeoutException("Condition not reached within " + timeout);
    }
}