natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.Tests/JetStream/JetStreamFileStoreReadQueryTests.cs

using System.Reflection;
using Shouldly;
using ZB.MOM.NatsNet.Server;
using ZB.MOM.NatsNet.Server.Internal.DataStructures;

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

public sealed class JetStreamFileStoreReadQueryTests
{
    [Theory]
    [InlineData("")]
    [InlineData(">")]
    public void CheckSkipFirstBlock_FilterIsAll_ReturnsNextBlock(string filter)
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(fs, NewBlock(1, 1, 10), NewBlock(2, 11, 20));

            var (next, error) = InvokeCheckSkipFirstBlock(fs, filter, wc: true, bi: 0);

            error.ShouldBeNull();
            next.ShouldBe(1);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void CheckSkipFirstBlock_LiteralFilterWithoutPsiMatch_ReturnsStoreEof()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(fs, NewBlock(1, 1, 10), NewBlock(2, 11, 20));
            SetPsims(fs, ("bar", 1, 2, 2));

            var (next, error) = InvokeCheckSkipFirstBlock(fs, "foo", wc: false, bi: 0);

            next.ShouldBe(-1);
            error.ShouldBe(StoreErrors.ErrStoreEOF);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void SelectSkipFirstBlock_StopAtCurrentBlock_ReturnsStoreEof()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(fs, NewBlock(1, 1, 10), NewBlock(3, 11, 20));

            var (next, error) = InvokeSelectSkipFirstBlock(fs, bi: 1, start: 1, stop: 3);

            next.ShouldBe(-1);
            error.ShouldBe(StoreErrors.ErrStoreEOF);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void SelectSkipFirstBlock_StartAfterCurrentBlock_ReturnsSelectedBlock()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(
                fs,
                NewBlock(1, 1, 10),
                NewBlock(5, 11, 20),
                NewBlock(9, 21, 30));

            var (next, error) = InvokeSelectSkipFirstBlock(fs, bi: 0, start: 5, stop: 9);

            error.ShouldBeNull();
            next.ShouldBe(1);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void CheckSkipFirstBlock_StopBeforeOrAtCurrentBlock_ReturnsStoreEof()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(fs, NewBlock(1, 1, 10), NewBlock(2, 11, 20), NewBlock(3, 21, 30));
            SetPsims(fs, ("foo", 1, 3, 3));

            var (next, error) = InvokeCheckSkipFirstBlock(fs, "foo", wc: false, bi: 2);

            next.ShouldBe(-1);
            error.ShouldBe(StoreErrors.ErrStoreEOF);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void CheckSkipFirstBlockMulti_IntersectingSubjectsOnly_SelectsMatchingBlock()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            ConfigureBlocks(
                fs,
                NewBlock(1, 1, 10),
                NewBlock(2, 11, 20),
                NewBlock(5, 21, 30),
                NewBlock(9, 31, 40));
            SetPsims(
                fs,
                ("zoo.c", 2, 2, 1),
                ("foo.a", 5, 5, 1),
                ("bar.b", 9, 9, 1));

            var sl = GenericSublist<EmptyStruct>.NewSimpleSublist();
            sl.Insert("foo.*", EmptyStruct.Value);

            var (next, error) = InvokeCheckSkipFirstBlockMulti(fs, sl, bi: 0);

            error.ShouldBeNull();
            next.ShouldBe(2);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Theory]
    [InlineData("")]
    [InlineData(">")]
    public void NumFilteredPending_FilterIsAll_UsesStreamState(string filter)
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            SetField(fs, "_state", new StreamState { FirstSeq = 5UL, LastSeq = 99UL, Msgs = 42UL });
            var ss = new SimpleState();

            InvokeNumFilteredPending(fs, filter, ss);

            ss.Msgs.ShouldBe(42UL);
            ss.First.ShouldBe(5UL);
            ss.Last.ShouldBe(99UL);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void NumFilteredPending_NoMatch_ResetsSimpleStateToZero()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            SetPsims(fs, ("foo.a", 1, 1, 1));
            var ss = new SimpleState { Msgs = 999UL, First = 123UL, Last = 456UL };

            InvokeNumFilteredPending(fs, "bar.b", ss);

            ss.Msgs.ShouldBe(0UL);
            ss.First.ShouldBe(0UL);
            ss.Last.ShouldBe(0UL);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void NumFilteredPending_LiteralAndWildcard_UsesPsiTotalsAndBlockBounds()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);

            var b1 = NewBlock(1, 1, 20);
            b1.Fss = BuildSubjectStateTree(("foo.a", 2, 10, 11));

            var b2 = NewBlock(2, 21, 40);
            b2.Fss = BuildSubjectStateTree(("foo.a", 1, 30, 30), ("foo.b", 3, 35, 40));

            ConfigureBlocks(fs, b1, b2);
            SetPsims(fs, ("foo.a", 1, 2, 3), ("foo.b", 2, 2, 3));

            var literal = new SimpleState();
            InvokeNumFilteredPending(fs, "foo.a", literal);
            literal.Msgs.ShouldBe(3UL);
            literal.First.ShouldBe(10UL);
            literal.Last.ShouldBe(30UL);

            var wildcard = new SimpleState();
            InvokeNumFilteredPending(fs, "foo.*", wildcard);
            wildcard.Msgs.ShouldBe(6UL);
            wildcard.First.ShouldBe(10UL);
            wildcard.Last.ShouldBe(40UL);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void NumFilteredPendingNoLast_LiteralMatch_LeavesLastAsZero()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            var b1 = NewBlock(1, 1, 20);
            b1.Fss = BuildSubjectStateTree(("foo.a", 2, 10, 11));
            ConfigureBlocks(fs, b1);
            SetPsims(fs, ("foo.a", 1, 1, 2));

            var ss = new SimpleState();
            InvokeNumFilteredPendingNoLast(fs, "foo.a", ss);

            ss.Msgs.ShouldBe(2UL);
            ss.First.ShouldBe(10UL);
            ss.Last.ShouldBe(0UL);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void NumFilteredPending_StaleFirstBlockHint_UpdatesPsiFirstBlock()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);

            var b1 = NewBlock(1, 1, 10);
            b1.Fss = BuildSubjectStateTree(("bar.a", 1, 1, 1));

            var b2 = NewBlock(2, 11, 20);
            b2.Fss = BuildSubjectStateTree(("foo.a", 1, 20, 20));

            ConfigureBlocks(fs, b1, b2);
            SetPsims(fs, ("foo.a", 1, 2, 1));
            var ss = new SimpleState();

            InvokeNumFilteredPending(fs, "foo.a", ss);

            ss.Msgs.ShouldBe(1UL);
            ss.First.ShouldBe(20UL);
            ss.Last.ShouldBe(20UL);

            var updated = false;
            for (var i = 0; i < 100; i++)
            {
                var psi = GetPsi(fs, "foo.a");
                if (psi != null && psi.Fblk == 2)
                {
                    updated = true;
                    break;
                }
                Thread.Sleep(10);
            }

            updated.ShouldBeTrue();
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void AllLastSeqsLocked_MultipleSubjects_ReturnsSortedLastSequences()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            SetField(fs, "_state", new StreamState { Msgs = 6UL });

            var b1 = NewBlock(1, 1, 20);
            b1.Fss = BuildSubjectStateTree(("foo.a", 2, 10, 10), ("foo.b", 2, 20, 20));
            var b2 = NewBlock(2, 21, 40);
            b2.Fss = BuildSubjectStateTree(("foo.b", 1, 25, 25), ("foo.c", 1, 15, 15));

            ConfigureBlocks(fs, b1, b2);
            SetPsims(fs, ("foo.a", 1, 1, 2), ("foo.b", 1, 2, 3), ("foo.c", 2, 2, 1));

            var (seqs, error) = InvokeAllLastSeqsLocked(fs);

            error.ShouldBeNull();
            seqs.ShouldBe([10UL, 15UL, 25UL]);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void AllLastSeqsLocked_NoMessagesOrNoTracking_ReturnsEmpty()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            SetField(fs, "_state", new StreamState { Msgs = 0UL });
            SetPsims(fs, ("foo.a", 1, 1, 1));

            var (noMsgsSeqs, noMsgsError) = InvokeAllLastSeqsLocked(fs);
            noMsgsError.ShouldBeNull();
            noMsgsSeqs.ShouldBeEmpty();
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }

        var noTrackDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(noTrackDir, subjects: []);
            SetField(fs, "_state", new StreamState { Msgs = 1UL });
            SetField(fs, "_psim", new SubjectTree<Psi>());

            var (noTrackSeqs, noTrackError) = InvokeAllLastSeqsLocked(fs);
            noTrackError.ShouldBeNull();
            noTrackSeqs.ShouldBeEmpty();
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(noTrackDir);
        }
    }

    [Fact]
    public void AllLastSeqsLocked_LastNeedsUpdate_RecalculatesBeforeCollecting()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir);
            SetField(fs, "_state", new StreamState { Msgs = 3UL });
            SetPsims(fs, ("foo.a", 1, 2, 3));

            var b1 = NewBlock(1, 1, 12);
            b1.Fss = BuildSubjectStateTree(("foo.a", 3, 1, 12));

            var b2 = NewBlock(2, 13, 20);
            var stale = new SimpleState { Msgs = 3UL, First = 1UL, Last = 0UL, LastNeedsUpdate = true };
            var staleTree = new SubjectTree<SimpleState>();
            staleTree.Insert(System.Text.Encoding.UTF8.GetBytes("foo.a"), stale);
            b2.Fss = staleTree;

            ConfigureBlocks(fs, b1, b2);

            var (seqs, error) = InvokeAllLastSeqsLocked(fs);

            error.ShouldBeNull();
            seqs.ShouldBe([12UL]);
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void FilterIsAll_ReorderedEquivalentFilters_ReturnsTrue()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir, ["foo.*", "bar.>"]);

            InvokeFilterIsAll(fs, ["bar.>", "foo.*"]).ShouldBeTrue();
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    [Fact]
    public void FilterIsAll_CountMismatchOrNonSubset_ReturnsFalse()
    {
        var storeDir = CreateStoreDir();
        try
        {
            var fs = CreateStore(storeDir, ["foo.*", "bar.>"]);

            InvokeFilterIsAll(fs, ["foo.A"]).ShouldBeFalse();
            InvokeFilterIsAll(fs, ["bar.>", "baz.*"]).ShouldBeFalse();
            fs.Stop();
        }
        finally
        {
            DeleteStoreDir(storeDir);
        }
    }

    private static string CreateStoreDir()
    {
        var root = Path.Combine(Path.GetTempPath(), $"fs-read-query-{Guid.NewGuid():N}");
        Directory.CreateDirectory(root);
        return root;
    }

    private static void DeleteStoreDir(string storeDir)
    {
        if (Directory.Exists(storeDir))
            Directory.Delete(storeDir, recursive: true);
    }

    private static JetStreamFileStore CreateStore(string storeDir, string[]? subjects = null)
    {
        subjects ??= ["foo.*", "bar.*", "zoo.*"];

        return new JetStreamFileStore(
            new FileStoreConfig { StoreDir = storeDir, BlockSize = 1024 },
            new FileStreamInfo
            {
                Created = DateTime.UtcNow,
                Config = new StreamConfig
                {
                    Name = "S",
                    Storage = StorageType.FileStorage,
                    Subjects = subjects,
                },
            });
    }

    private static MessageBlock NewBlock(uint index, ulong first, ulong last)
    {
        return new MessageBlock
        {
            Index = index,
            First = new MsgId { Seq = first },
            Last = new MsgId { Seq = last },
        };
    }

    private static SubjectTree<SimpleState> BuildSubjectStateTree(params (string Subject, ulong Msgs, ulong First, ulong Last)[] states)
    {
        var tree = new SubjectTree<SimpleState>();
        foreach (var (subject, msgs, first, last) in states)
        {
            tree.Insert(
                System.Text.Encoding.UTF8.GetBytes(subject),
                new SimpleState
                {
                    Msgs = msgs,
                    First = first,
                    Last = last,
                });
        }

        return tree;
    }

    private static void ConfigureBlocks(JetStreamFileStore fs, params MessageBlock[] blocks)
    {
        var ordered = blocks.OrderBy(b => b.Index).ToList();
        var bim = ordered.ToDictionary(b => b.Index, b => b);
        SetField(fs, "_blks", ordered);
        SetField(fs, "_bim", bim);
        SetField(fs, "_lmb", ordered.Count > 0 ? ordered[^1] : null);
    }

    private static void SetPsims(JetStreamFileStore fs, params (string Subject, uint Fblk, uint Lblk, ulong Total)[] entries)
    {
        var psim = new SubjectTree<Psi>();
        foreach (var (subject, fblk, lblk, total) in entries)
        {
            psim.Insert(
                System.Text.Encoding.UTF8.GetBytes(subject),
                new Psi
                {
                    Fblk = fblk,
                    Lblk = lblk,
                    Total = total,
                });
        }

        SetField(fs, "_psim", psim);
    }

    private static (int Next, Exception? Error) InvokeCheckSkipFirstBlock(JetStreamFileStore fs, string filter, bool wc, int bi)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("CheckSkipFirstBlock", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        var result = mi!.Invoke(fs, [filter, wc, bi]);
        result.ShouldNotBeNull();
        return ((int, Exception?))result!;
    }

    private static (int Next, Exception? Error) InvokeCheckSkipFirstBlockMulti(JetStreamFileStore fs, SimpleSublist sl, int bi)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("CheckSkipFirstBlockMulti", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        var result = mi!.Invoke(fs, [sl, bi]);
        result.ShouldNotBeNull();
        return ((int, Exception?))result!;
    }

    private static (int Next, Exception? Error) InvokeSelectSkipFirstBlock(JetStreamFileStore fs, int bi, uint start, uint stop)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("SelectSkipFirstBlock", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        var result = mi!.Invoke(fs, [bi, start, stop]);
        result.ShouldNotBeNull();
        return ((int, Exception?))result!;
    }

    private static void InvokeNumFilteredPending(JetStreamFileStore fs, string filter, SimpleState ss)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("NumFilteredPending", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        _ = mi!.Invoke(fs, [filter, ss]);
    }

    private static void InvokeNumFilteredPendingNoLast(JetStreamFileStore fs, string filter, SimpleState ss)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("NumFilteredPendingNoLast", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        _ = mi!.Invoke(fs, [filter, ss]);
    }

    private static (ulong[] Seqs, Exception? Error) InvokeAllLastSeqsLocked(JetStreamFileStore fs)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("AllLastSeqsLocked", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        var result = mi!.Invoke(fs, []);
        result.ShouldNotBeNull();
        return ((ulong[], Exception?))result!;
    }

    private static bool InvokeFilterIsAll(JetStreamFileStore fs, string[] filters)
    {
        var mi = typeof(JetStreamFileStore).GetMethod("FilterIsAll", BindingFlags.Instance | BindingFlags.NonPublic);
        mi.ShouldNotBeNull();
        var result = mi!.Invoke(fs, [filters]);
        result.ShouldNotBeNull();
        return (bool)result;
    }

    private static Psi? GetPsi(JetStreamFileStore fs, string subject)
    {
        var fi = typeof(JetStreamFileStore).GetField("_psim", BindingFlags.Instance | BindingFlags.NonPublic);
        fi.ShouldNotBeNull();
        var psim = fi!.GetValue(fs).ShouldBeOfType<SubjectTree<Psi>>();
        var (psi, found) = psim.Find(System.Text.Encoding.UTF8.GetBytes(subject));
        return found ? psi : null;
    }

    private static void SetField<T>(object target, string fieldName, T value)
    {
        var fi = target.GetType().GetField(fieldName, BindingFlags.Instance | BindingFlags.NonPublic);
        fi.ShouldNotBeNull();
        fi!.SetValue(target, value);
    }
}