natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.Tests/ImplBacklog/ConcurrencyTests1.Impltests.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using Shouldly;
using ZB.MOM.NatsNet.Server;

namespace ZB.MOM.NatsNet.Server.Tests.ImplBacklog;

public sealed partial class ConcurrencyTests1
{
    [Fact] // T:2390
    public void NoRaceJetStreamClusterLargeStreamInlineCatchup_ShouldSucceed()
    {
        var cluster = new JetStreamCluster
        {
            Streams = new Dictionary<string, Dictionary<string, StreamAssignment>>
            {
                ["A"] = new Dictionary<string, StreamAssignment>
                {
                    ["BIG"] = new()
                    {
                        Config = new StreamConfig { Name = "BIG", Subjects = ["big.>"] },
                    },
                },
            },
        };

        var engine = new JetStreamEngine(new global::ZB.MOM.NatsNet.Server.JetStream { Cluster = cluster });
        engine.SubjectsOverlap("A", ["big.orders"]).ShouldBeTrue();
    }

    [Fact] // T:2459
    public void NoRaceJetStreamClusterDifferentRTTInterestBasedStreamSetup_ShouldSucceed()
    {
        var updates = new RecoveryUpdates();
        var stream = new StreamAssignment { Client = new ClientInfo { Account = "A" }, Config = new StreamConfig { Name = "RTT" } };

        updates.AddStream(stream);
        updates.UpdateStream(stream);
        updates.UpdateStreams.ShouldContainKey("A:RTT");
    }

    [Fact] // T:2422
    public void NoRaceJetStreamConsumerFileStoreConcurrentDiskIO_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            const int consumerCount = 400;
            var start = new ManualResetEventSlim(false);
            var errors = new ConcurrentQueue<Exception>();
            var timestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds() * 1_000_000_000L;

            var workers = new List<Task>(consumerCount);
            for (var i = 0; i < consumerCount; i++)
            {
                var consumer = fs.ConsumerStore(
                    $"o{i}",
                    DateTime.UtcNow,
                    new ConsumerConfig { AckPolicy = AckPolicy.AckExplicit });

                workers.Add(Task.Run(() =>
                {
                    try
                    {
                        start.Wait(TimeSpan.FromSeconds(5));
                        consumer.UpdateDelivered(22, 22, 1, timestamp);
                        consumer.EncodedState();
                        consumer.Delete();
                    }
                    catch (Exception ex)
                    {
                        errors.Enqueue(ex);
                    }
                }));
            }

            start.Set();
            Task.WaitAll(workers.ToArray());
            errors.ShouldBeEmpty();
        });
    }

    [Fact] // T:2452
    public void NoRaceFileStoreStreamMaxAgePerformance_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            Parallel.For(0, 200, i => fs.StoreMsg($"age.{i % 4}", null, new[] { (byte)(i % 255) }, 0));

            var state = fs.State();
            state.Msgs.ShouldBeGreaterThan(0UL);
            state.LastSeq.ShouldBeGreaterThanOrEqualTo(state.Msgs);

            var (total, validThrough, err) = fs.NumPending(1, ">", false);
            err.ShouldBeNull();
            total.ShouldBeGreaterThan(0UL);
            validThrough.ShouldBeGreaterThan(0UL);
        }, DefaultStreamConfig(maxAge: TimeSpan.FromMilliseconds(20)));
    }

    [Fact] // T:2453
    public void NoRaceFileStoreFilteredStateWithLargeDeletes_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            for (var i = 0; i < 240; i++)
                fs.StoreMsg("fd", null, new[] { (byte)(i % 255) }, 0);

            Parallel.For(1L, 240L, i =>
            {
                if (i % 3 == 0)
                    fs.RemoveMsg((ulong)i);
            });

            var filtered = fs.FilteredState(1, "fd");
            filtered.Msgs.ShouldBeGreaterThan(0UL);
            filtered.Last.ShouldBeGreaterThanOrEqualTo(filtered.First);

            fs.SubjectsTotals(">")["fd"].ShouldBeGreaterThan(0UL);
        });
    }

    [Fact] // T:2462
    public void NoRaceFileStoreNumPending_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            for (var i = 0; i < 100; i++)
                fs.StoreMsg($"np.{i % 5}", null, "x"u8.ToArray(), 0);

            var errors = new ConcurrentQueue<Exception>();
            var workers = Enumerable.Range(0, 8).Select(_ => Task.Run(() =>
            {
                try
                {
                    for (var i = 0; i < 40; i++)
                    {
                        var (_, _, err1) = fs.NumPending(1, ">", false);
                        if (err1 != null)
                            throw err1;

                        var (_, _, err2) = fs.NumPendingMulti(1, new[] { "np.1", "np.*" }, false);
                        if (err2 != null)
                            throw err2;
                    }
                }
                catch (Exception ex)
                {
                    errors.Enqueue(ex);
                }
            })).ToArray();

            Task.WaitAll(workers);
            errors.ShouldBeEmpty();
        });
    }

    [Fact] // T:2427
    public void NoRaceJetStreamFileStoreKeyFileCleanup_ShouldSucceed()
    {
        WithStore((_, root) =>
        {
            var msgDir = Path.Combine(root, FileStoreDefaults.MsgDir);
            Directory.CreateDirectory(msgDir);
            var perm = UnixFileMode.UserRead | UnixFileMode.UserWrite;

            var errors = new ConcurrentQueue<Exception>();
            Parallel.For(0, 300, i =>
            {
                var payload = BitConverter.GetBytes(i);
                var keyFile = Path.Combine(msgDir, string.Format(FileStoreDefaults.KeyScan, (uint)(i + 1)));
                var err = JetStreamFileStore.WriteAtomically(keyFile, payload, perm, sync: true);
                if (err != null)
                    errors.Enqueue(err);
            });

            errors.ShouldBeEmpty();
            var keyFiles = Directory.GetFiles(msgDir, "*.key");
            keyFiles.Length.ShouldBe(300);

            foreach (var key in keyFiles.Skip(1))
                File.Delete(key);

            Directory.GetFiles(msgDir, "*.key").Length.ShouldBe(1);
            Directory.GetFiles(msgDir, "*.tmp").ShouldBeEmpty();
        });
    }

    [Fact] // T:2447
    public void NoRaceEncodeConsumerStateBug_ShouldSucceed()
    {
        for (var i = 0; i < 5_000; i++)
        {
            var pending = new Pending
            {
                Sequence = 1,
                Timestamp = DateTimeOffset.UtcNow.AddSeconds(1).ToUnixTimeSeconds() * 1_000_000_000L,
            };
            var state = new ConsumerState
            {
                Delivered = new SequencePair { Consumer = 1, Stream = 1 },
                Pending = new Dictionary<ulong, Pending> { [1] = pending },
            };

            var encoded = StoreParity.EncodeConsumerState(state);
            var (_, err) = JetStreamFileStore.DecodeConsumerState(encoded);
            err.ShouldBeNull();
        }
    }

    [Fact]
    public void NoRaceJetStreamConsumerDeleteWithFlushPending_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            const int consumerCount = 100;
            var errors = new ConcurrentQueue<Exception>();
            var ts = DateTimeOffset.UtcNow.ToUnixTimeSeconds() * 1_000_000_000L;
            var workers = new List<Task>(consumerCount);

            for (var i = 0; i < consumerCount; i++)
            {
                var consumer = fs.ConsumerStore(
                    $"flush-del-{i}",
                    DateTime.UtcNow,
                    new ConsumerConfig { AckPolicy = AckPolicy.AckExplicit });

                workers.Add(Task.Run(() =>
                {
                    var updater = Task.Run(() =>
                    {
                        for (ulong n = 1; n <= 50; n++)
                        {
                            try
                            {
                                consumer.UpdateDelivered(n, n, 1, ts + (long)n);
                            }
                            catch (Exception ex) when (ReferenceEquals(ex, StoreErrors.ErrStoreClosed))
                            {
                                break;
                            }
                        }
                    });

                    Thread.Sleep(1);

                    try
                    {
                        consumer.Delete();
                        updater.Wait();
                    }
                    catch (Exception ex)
                    {
                        errors.Enqueue(ex);
                    }
                }));
            }

            Task.WaitAll(workers.ToArray());
            errors.ShouldBeEmpty();
        });
    }

    [Fact] // T:2441
    public void NoRaceJetStreamFileStoreLargeKVAccessTiming_ShouldSucceed()
    {
        var value = Enumerable.Repeat((byte)'Z', 256).ToArray();
        const int keyCount = 5_000;

        WithStore((fs, _) =>
        {
            for (var i = 1; i <= keyCount; i++)
            {
                fs.StoreMsg($"KV.STREAM_NAME.{i}", null, value, 0).Seq.ShouldBeGreaterThan(0UL);
            }

            var sw = Stopwatch.StartNew();
            var last = fs.LoadLastMsg($"KV.STREAM_NAME.{keyCount}", null);
            sw.Stop();
            var lastLookup = sw.Elapsed;

            last.ShouldNotBeNull();
            last!.Msg.ShouldBe(value);

            sw.Restart();
            var first = fs.LoadLastMsg("KV.STREAM_NAME.1", null);
            sw.Stop();
            var firstLookup = sw.Elapsed;

            first.ShouldNotBeNull();
            first!.Msg.ShouldBe(value);

            // Keep generous bounds to avoid machine-specific flakiness while still
            // asserting access stays fast under a large key set.
            lastLookup.ShouldBeLessThan(TimeSpan.FromMilliseconds(250));
            firstLookup.ShouldBeLessThan(TimeSpan.FromMilliseconds(350));

            var firstState = fs.FilteredState(1, "KV.STREAM_NAME.1");
            var lastState = fs.FilteredState(1, $"KV.STREAM_NAME.{keyCount}");
            firstState.First.ShouldBeGreaterThan(0UL);
            lastState.First.ShouldBeGreaterThan(0UL);
        }, DefaultStreamConfig());
    }

    [Fact] // T:2371
    public void NoRaceAvoidSlowConsumerBigMessages_ShouldSucceed()
    {
        WithStore((fs, _) =>
        {
            var errors = new ConcurrentQueue<Exception>();
            var payload = new byte[128 * 1024];

            Parallel.For(0, 40, i =>
            {
                try
                {
                    fs.StoreMsg($"big.{i % 4}", null, payload, 0).Seq.ShouldBeGreaterThan(0UL);
                    var sm = fs.LoadLastMsg($"big.{i % 4}", null);
                    sm.ShouldNotBeNull();
                    sm!.Msg.Length.ShouldBe(payload.Length);
                }
                catch (Exception ex)
                {
                    errors.Enqueue(ex);
                }
            });

            errors.ShouldBeEmpty();
            fs.State().Msgs.ShouldBeGreaterThan(0UL);
        });
    }

    [Fact] // T:2384
    public void NoRaceAcceptLoopsDoNotLeaveOpenedConn_ShouldSucceed()
    {
        var errors = new ConcurrentQueue<Exception>();

        Parallel.For(0, 20, _ =>
        {
            TcpListener? listener = null;
            TcpClient? client = null;
            TcpClient? accepted = null;

            try
            {
                listener = new TcpListener(IPAddress.Loopback, 0);
                listener.Start();
                var endpoint = (IPEndPoint)listener.LocalEndpoint;

                var acceptTask = listener.AcceptTcpClientAsync();
                client = new TcpClient();
                client.Connect(endpoint.Address, endpoint.Port);
                accepted = acceptTask.GetAwaiter().GetResult();

                accepted.Connected.ShouldBeTrue();
            }
            catch (Exception ex)
            {
                errors.Enqueue(ex);
            }
            finally
            {
                accepted?.Close();
                client?.Close();
                listener?.Stop();
            }
        });

        errors.ShouldBeEmpty();
    }

    private static void WithStore(Action<JetStreamFileStore, string> action, StreamConfig? cfg = null)
    {
        var root = NewRoot();
        Directory.CreateDirectory(root);
        JetStreamFileStore? fs = null;

        try
        {
            fs = JetStreamFileStore.NewFileStore(new FileStoreConfig { StoreDir = root }, cfg ?? DefaultStreamConfig());
            action(fs, root);
        }
        finally
        {
            fs?.Stop();
            if (Directory.Exists(root))
                Directory.Delete(root, recursive: true);
        }
    }

    private static StreamConfig DefaultStreamConfig(TimeSpan? maxAge = null)
    {
        return new StreamConfig
        {
            Name = "TEST",
            Storage = StorageType.FileStorage,
            Subjects = ["test.>"],
            MaxMsgs = -1,
            MaxBytes = -1,
            MaxAge = maxAge ?? TimeSpan.Zero,
            MaxMsgsPer = -1,
            Discard = DiscardPolicy.DiscardOld,
            Retention = RetentionPolicy.LimitsPolicy,
        };
    }

    private static string NewRoot() => Path.Combine(Path.GetTempPath(), $"impl-fs-c1-{Guid.NewGuid():N}");
}