ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.IntegrationTests/Grpc/GrpcStreamBatchingIntegrationTests.cs

using System.Diagnostics;
using Akka.Actor;
using Akka.TestKit.Xunit2;
using Google.Protobuf;
using Grpc.Core;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using NSubstitute;
using Xunit.Abstractions;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Streaming;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;
using ZB.MOM.WW.ScadaBridge.SiteRuntime;
using ZB.MOM.WW.ScadaBridge.SiteRuntime.Streaming;

namespace ZB.MOM.WW.ScadaBridge.IntegrationTests.Grpc;

/// <summary>
/// End-to-end coverage for R2 — gRPC event batching on the site→central
/// <c>SiteStreamService</c> stream.
///
/// <para>
/// The chain assembled here is the real one, mocking only the HTTP/2 transport:
/// domain event → real <see cref="SiteStreamManager"/> broadcast → real
/// <see cref="SiteStreamGrpcServer"/> handler → real <c>StreamRelayActor</c> → real
/// coalescing pump → <b>proto serialize/parse round-trip</b> (what the wire actually
/// carries) → real <see cref="SiteStreamGrpcClient.ForEachEvent"/> unpack → real
/// <c>ConvertToDomainEvent</c>. The serialize/parse step is what makes these
/// version-skew claims real rather than in-memory object graph assertions.
/// </para>
///
/// <para>
/// <b>Version skew is covered in both directions.</b> Negotiation is a single additive
/// request field (<c>batching_supported</c>), whose proto3 default of false IS the
/// compatibility mechanism: an old central cannot set it, so a new site never sends it a
/// frame case its generated code would drop; a new central always sets it, and an old
/// site ignores the unknown field and keeps sending per-event frames the new client
/// accepts unchanged.
/// </para>
/// </summary>
public class GrpcStreamBatchingIntegrationTests(ITestOutputHelper output) : TestKit
{
    private const string Instance = "SiteA.Pump01";

    /// <summary>
    /// End-to-end latency threshold the target-scale load test asserts a P99 against
    /// (measured P99 there: 4.57 ms). The coalescing window is the only latency batching
    /// can add, so the batched pipe must stay comfortably inside the same budget.
    /// </summary>
    private static readonly TimeSpan LatencyThreshold = TimeSpan.FromMilliseconds(250);

    // ── Round trip: batched frames deliver every event, in order ────────────────

    [Fact]
    public async Task NegotiatedStream_RoundTripsEveryEventThroughTheWire_InOrder()
    {
        var (server, manager, frames, cts, streamTask) = await StartAsync(batchingSupported: true);

        const int total = 600;
        var t0 = new DateTimeOffset(2026, 8, 15, 12, 0, 0, TimeSpan.Zero);
        for (var i = 0; i < total; i++)
        {
            manager.PublishAttributeValueChanged(new AttributeValueChanged(
                Instance, "Modules.IO", "Seq", i, "Good", t0.AddMilliseconds(i)));
        }

        await WaitForConditionAsync(() => TotalEvents(frames) >= total, 30_000);
        cts.Cancel();
        await streamTask;

        var wire = SnapshotThroughTheWire(frames);

        // Batching actually happened — otherwise this test proves nothing about batching.
        Assert.Contains(wire, f => f.EventCase == SiteStreamEvent.EventOneofCase.Batch);
        Assert.True(wire.Count < total,
            $"{wire.Count} frames for {total} events — no coalescing happened");

        var delivered = Unpack(wire);
        Assert.Equal(total, delivered.Count);

        // Every event, exactly once, in the order the site produced it — and with its OWN
        // timestamp, not a frame-level one (end-to-end latency measurement rides it).
        Assert.Equal(
            Enumerable.Range(0, total).Select(i => i.ToString()),
            delivered.Select(e => e.Value));
        Assert.Equal(
            Enumerable.Range(0, total).Select(i => t0.AddMilliseconds(i)),
            delivered.Select(e => e.Timestamp));

        output.WriteLine($"round-trip: {total} events in {wire.Count} frames " +
            $"(mean {(double)total / wire.Count:0.0} events/frame)");

        GC.KeepAlive(server);
    }

    // ── Latency cost of the default window ─────────────────────────────────────

    [Fact]
    public async Task DefaultWindow_KeepsPerEventLatencyFarBelowTheThreshold()
    {
        // A trickle-with-backlog workload is the case the coalescing window actually
        // bites on: each burst is far short of the 100-event size cap, so the batch is
        // closed by the 25 ms window rather than by the cap. That makes this the WORST
        // case for added latency, not the best.
        var options = new CommunicationOptions();
        var (server, manager, frames, cts, streamTask) = await StartAsync(
            batchingSupported: true, options: options);

        const int bursts = 150;
        const int perBurst = 8;
        var stamps = new Dictionary<int, DateTimeOffset>();

        var seq = 0;
        for (var b = 0; b < bursts; b++)
        {
            for (var i = 0; i < perBurst; i++)
            {
                var ts = DateTimeOffset.UtcNow;
                stamps[seq] = ts;
                manager.PublishAttributeValueChanged(new AttributeValueChanged(
                    Instance, "Modules.IO", "Seq", seq, "Good", ts));
                seq++;
            }
            await Task.Delay(5);
        }

        var total = seq;
        await WaitForConditionAsync(() => TotalEvents(frames) >= total, 60_000);
        cts.Cancel();
        await streamTask;

        // Latency = the event's own site-side timestamp → the instant the frame carrying
        // it was handed to the response stream. That interval contains the coalescing
        // window and nothing else the pre-batching pipe did not already have.
        var latencies = new List<double>();
        lock (frames)
        {
            foreach (var (frame, writtenAt) in frames)
            {
                foreach (var evt in Flatten(frame))
                {
                    var s = int.Parse(evt.AttributeChanged.Value);
                    latencies.Add((writtenAt - stamps[s]).TotalMilliseconds);
                }
            }
        }

        latencies.Sort();
        var p50 = latencies[(int)(latencies.Count * 0.50)];
        var p99 = latencies[(int)(latencies.Count * 0.99)];
        var max = latencies[^1];

        output.WriteLine(
            $"window={options.GrpcStreamBatchWindow.TotalMilliseconds:0} ms cap={options.GrpcStreamBatchMaxEvents} " +
            $"events={latencies.Count} P50={p50:0.00} ms P99={p99:0.00} ms max={max:0.00} ms");

        Assert.Equal(total, latencies.Count);
        Assert.True(p99 < LatencyThreshold.TotalMilliseconds,
            $"P99 {p99:0.00} ms exceeded the {LatencyThreshold.TotalMilliseconds:0} ms end-to-end threshold");

        GC.KeepAlive(server);
    }

    // ── Version skew: OLD central ↔ NEW site ───────────────────────────────────

    [Fact]
    public async Task OldCentral_AgainstNewSite_NeverReceivesABatchFrame()
    {
        // An old central's InstanceStreamRequest bytes simply have no field 3 — build
        // exactly those bytes and let the NEW site parse them, so the negotiation default
        // is exercised off the wire rather than asserted on an object.
        var oldCentralBytes = BuildLegacyInstanceRequest("corr-old-central", Instance);
        var request = InstanceStreamRequest.Parser.ParseFrom(oldCentralBytes);
        Assert.False(request.BatchingSupported);

        var (server, manager, frames, cts, streamTask) = await StartAsync(request);

        const int total = 300;
        for (var i = 0; i < total; i++)
        {
            manager.PublishAttributeValueChanged(new AttributeValueChanged(
                Instance, "Modules.IO", "Seq", i, "Good", DateTimeOffset.UtcNow));
        }

        await WaitForConditionAsync(() => TotalEvents(frames) >= total, 30_000);
        cts.Cancel();
        await streamTask;

        var wire = SnapshotThroughTheWire(frames);

        // One event per frame, and — checked at the byte level, since that is what the
        // old peer's parser sees — never the field-4 batch tag.
        Assert.Equal(total, wire.Count);
        Assert.All(wire, f =>
            Assert.Equal(SiteStreamEvent.EventOneofCase.AttributeChanged, f.EventCase));
        Assert.All(wire, f => Assert.DoesNotContain(4, FieldNumbers(f)));

        GC.KeepAlive(server);
    }

    [Fact]
    public void BatchFrameRidesFieldFour_WhichAPreBatchingParserWouldDropSilently()
    {
        // WHY negotiation is mandatory rather than "just send batches". A batch frame is a
        // length-delimited field 4: an older generated parser skips it into unknown fields
        // and reports EventOneofCase.None, whose ConvertToDomainEvent returns null — the
        // whole batch would vanish with no error anywhere. The proto3 default on
        // batching_supported is what guarantees such a peer never receives one.
        var batch = new SiteStreamEvent
        {
            CorrelationId = "corr-shape",
            Batch = new SiteStreamEventBatch
            {
                Events = { MakeAttributeEvent(1), MakeAttributeEvent(2) }
            }
        };

        var fields = FieldNumbers(batch);
        Assert.Contains(4, fields);
        Assert.DoesNotContain(2, fields);
        Assert.DoesNotContain(3, fields);

        // Field 4 is length-delimited (wire type 2) — the shape an unknown-field-tolerant
        // parser can skip without corrupting the rest of the message.
        Assert.Equal(2u, WireTypeOfField(batch, 4));

        // And the per-event frames a pre-batching site emits still parse and convert on the
        // NEW client (the other skew direction, at the same byte level).
        var plain = SiteStreamEvent.Parser.ParseFrom(MakeAttributeEvent(7).ToByteArray());
        Assert.NotNull(SiteStreamGrpcClient.ConvertToDomainEvent(plain));
    }

    // ── Version skew: NEW central ↔ OLD site ───────────────────────────────────

    [Fact]
    public async Task NewCentral_AgainstOldSite_StillReceivesEveryEvent()
    {
        // An old site ignores batching_supported and emits per-event frames. That emission
        // shape is exactly what the current server produces with batching off, so drive
        // the real server that way and feed the result through the NEW client's unpack —
        // which must handle the single-event case identically to before R2.
        var (server, manager, frames, cts, streamTask) = await StartAsync(batchingSupported: false);

        const int total = 200;
        var t0 = new DateTimeOffset(2026, 8, 15, 13, 0, 0, TimeSpan.Zero);
        for (var i = 0; i < total; i++)
        {
            manager.PublishAttributeValueChanged(new AttributeValueChanged(
                Instance, "Modules.IO", "Seq", i, "Good", t0.AddMilliseconds(i)));
        }

        await WaitForConditionAsync(() => TotalEvents(frames) >= total, 30_000);
        cts.Cancel();
        await streamTask;

        var wire = SnapshotThroughTheWire(frames);
        Assert.Equal(total, wire.Count);

        var delivered = Unpack(wire);
        Assert.Equal(
            Enumerable.Range(0, total).Select(i => i.ToString()),
            delivered.Select(e => e.Value));
        Assert.Equal(
            Enumerable.Range(0, total).Select(i => t0.AddMilliseconds(i)),
            delivered.Select(e => e.Timestamp));

        GC.KeepAlive(server);
    }

    // ── Helpers ────────────────────────────────────────────────────────────────

    private static SiteStreamEvent MakeAttributeEvent(int seq) => new()
    {
        CorrelationId = "corr-shape",
        AttributeChanged = new AttributeValueUpdate
        {
            InstanceUniqueName = Instance,
            AttributePath = "Modules.IO",
            AttributeName = "Seq",
            Value = seq.ToString(),
            Quality = Quality.Good,
            Timestamp = Google.Protobuf.WellKnownTypes.Timestamp.FromDateTimeOffset(DateTimeOffset.UnixEpoch)
        }
    };

    /// <summary>
    /// Serializes an <c>InstanceStreamRequest</c> the way a central built BEFORE R2 would:
    /// fields 1 and 2 only, with no <c>batching_supported</c> on the wire at all.
    /// </summary>
    private static byte[] BuildLegacyInstanceRequest(string correlationId, string instance)
    {
        using var ms = new MemoryStream();
        var output = new CodedOutputStream(ms);
        output.WriteTag(1, WireFormat.WireType.LengthDelimited);
        output.WriteString(correlationId);
        output.WriteTag(2, WireFormat.WireType.LengthDelimited);
        output.WriteString(instance);
        output.Flush();
        return ms.ToArray();
    }

    /// <summary>Top-level field numbers present in a serialized message.</summary>
    private static HashSet<int> FieldNumbers(IMessage message)
    {
        var fields = new HashSet<int>();
        var input = new CodedInputStream(message.ToByteArray());
        uint tag;
        while ((tag = input.ReadTag()) != 0)
        {
            fields.Add(WireFormat.GetTagFieldNumber(tag));
            input.SkipLastField();
        }
        return fields;
    }

    /// <summary>Wire type of the given top-level field number in a serialized message.</summary>
    private static uint WireTypeOfField(IMessage message, int fieldNumber)
    {
        var input = new CodedInputStream(message.ToByteArray());
        uint tag;
        while ((tag = input.ReadTag()) != 0)
        {
            if (WireFormat.GetTagFieldNumber(tag) == fieldNumber)
                return (uint)WireFormat.GetTagWireType(tag);
            input.SkipLastField();
        }
        throw new InvalidOperationException($"field {fieldNumber} not present");
    }

    private static IEnumerable<SiteStreamEvent> Flatten(SiteStreamEvent frame)
    {
        if (frame.EventCase == SiteStreamEvent.EventOneofCase.Batch)
        {
            foreach (var inner in frame.Batch.Events) yield return inner;
            yield break;
        }
        yield return frame;
    }

    private static int TotalEvents(List<(SiteStreamEvent Frame, DateTimeOffset WrittenAt)> frames)
    {
        lock (frames) { return frames.Sum(f => Flatten(f.Frame).Count()); }
    }

    /// <summary>
    /// Takes the captured frames through a real protobuf serialize/parse round-trip — the
    /// step that makes every claim in this file about wire compatibility a wire claim.
    /// </summary>
    private static List<SiteStreamEvent> SnapshotThroughTheWire(
        List<(SiteStreamEvent Frame, DateTimeOffset WrittenAt)> frames)
    {
        lock (frames)
        {
            return [.. frames.Select(f => SiteStreamEvent.Parser.ParseFrom(f.Frame.ToByteArray()))];
        }
    }

    /// <summary>Unpacks wire frames through the REAL client path into domain events.</summary>
    private static List<AttributeValueChanged> Unpack(IEnumerable<SiteStreamEvent> wire)
    {
        var delivered = new List<AttributeValueChanged>();
        foreach (var frame in wire)
        {
            SiteStreamGrpcClient.ForEachEvent(frame, e =>
            {
                if (SiteStreamGrpcClient.ConvertToDomainEvent(e) is AttributeValueChanged a)
                    delivered.Add(a);
            });
        }
        return delivered;
    }

    private Task<(SiteStreamGrpcServer Server, SiteStreamManager Manager,
        List<(SiteStreamEvent Frame, DateTimeOffset WrittenAt)> Frames,
        CancellationTokenSource Cts, Task StreamTask)>
        StartAsync(bool batchingSupported, CommunicationOptions? options = null)
        => StartAsync(new InstanceStreamRequest
        {
            CorrelationId = "corr-batching",
            InstanceUniqueName = Instance,
            BatchingSupported = batchingSupported
        }, options);

    /// <summary>
    /// Brings up a real site broadcast hub + real gRPC server handler for the supplied
    /// subscription request, capturing every written frame with the instant it was written.
    /// </summary>
    private async Task<(SiteStreamGrpcServer Server, SiteStreamManager Manager,
        List<(SiteStreamEvent Frame, DateTimeOffset WrittenAt)> Frames,
        CancellationTokenSource Cts, Task StreamTask)>
        StartAsync(InstanceStreamRequest request, CommunicationOptions? options = null)
    {
        var manager = new SiteStreamManager(
            new SiteRuntimeOptions { StreamBufferSize = 4096 },
            NullLogger<SiteStreamManager>.Instance);
        manager.Initialize(Sys);

        var server = new SiteStreamGrpcServer(
            manager,
            NullLogger<SiteStreamGrpcServer>.Instance,
            Options.Create(options ?? new CommunicationOptions()));
        server.SetReady(Sys);

        var frames = new List<(SiteStreamEvent, DateTimeOffset)>();
        var writer = Substitute.For<IServerStreamWriter<SiteStreamEvent>>();
        writer.WriteAsync(Arg.Any<SiteStreamEvent>(), Arg.Any<CancellationToken>())
            .Returns(Task.CompletedTask)
            .AndDoes(ci =>
            {
                var frame = ci.Arg<SiteStreamEvent>();
                var at = DateTimeOffset.UtcNow;
                lock (frames) { frames.Add((frame, at)); }
            });

        var cts = new CancellationTokenSource();
        var context = Substitute.For<ServerCallContext>();
        context.CancellationToken.Returns(cts.Token);

        var streamTask = Task.Run(() => server.SubscribeInstance(request, writer, context));

        // The publish must not race the materialized subscription.
        await WaitForConditionAsync(() => manager.SubscriptionCount == 1);

        return (server, manager, frames, cts, streamTask);
    }

    private static async Task WaitForConditionAsync(Func<bool> condition, int timeoutMs = 5000)
    {
        var started = Stopwatch.GetTimestamp();
        while (!condition() && Stopwatch.GetElapsedTime(started) < TimeSpan.FromMilliseconds(timeoutMs))
        {
            await Task.Delay(10);
        }
        Assert.True(condition(), $"Condition not met within {timeoutMs}ms");
    }
}