natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Storage/FileStoreTtlTests.cs

// Reference: golang/nats-server/server/filestore_test.go
// Tests ported:
//   TestFileStoreWriteCache               — write cache hit (msgBlock.cache)
//   TestFileStoreClearCache               — ClearCache evicts, disk read still works
//   TestFileStoreTtlWheelExpiry           — TTL wheel expires old messages (expireMsgs)
//   TestFileStoreTtlWheelRetention        — TTL wheel retains unexpired messages
//   TestFileStoreStoreMsg                 — StoreMsg returns seq + timestamp
//   TestFileStoreStoreMsgPerMsgTtl        — StoreMsg with per-message TTL
//   TestFileStoreRecoveryReregiistersTtls — recovery re-registers unexpired TTL entries

using System.Text;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.JetStream.Tests.JetStream.Storage;

/// <summary>
/// Tests for the MsgBlock write cache and FileStore TTL wheel scheduling.
/// Reference: golang/nats-server/server/filestore.go — msgBlock.cache, expireMsgs, storeMsg TTL.
/// </summary>
public sealed class FileStoreTtlTests : IDisposable
{
    private readonly string _dir;

    public FileStoreTtlTests()
    {
        _dir = Path.Combine(Path.GetTempPath(), $"nats-js-ttl-{Guid.NewGuid():N}");
        Directory.CreateDirectory(_dir);
    }

    public void Dispose()
    {
        if (Directory.Exists(_dir))
            Directory.Delete(_dir, recursive: true);
    }

    private FileStore CreateStore(FileStoreOptions? options = null, string? sub = null)
    {
        var dir = sub is null ? _dir : Path.Combine(_dir, sub);
        var opts = options ?? new FileStoreOptions();
        opts.Directory = dir;
        return new FileStore(opts);
    }

    // -------------------------------------------------------------------------
    // MsgBlock write cache tests
    // -------------------------------------------------------------------------

    // Go: TestFileStoreWriteCache — filestore_test.go (msgBlock.cache hit path)
    [Fact]
    public async Task WriteCache_ReadReturnsFromCache()
    {
        // The active block maintains a write cache populated on every Write/WriteAt.
        // After writing a message, the active block's cache should contain it so
        // Read() returns without touching disk.
        await using var store = CreateStore();

        var seq = await store.AppendAsync("foo", "hello"u8.ToArray(), default);
        seq.ShouldBe(1UL);

        // Load back through the store's in-memory cache (which calls MsgBlock.Read internally).
        var msg = await store.LoadAsync(seq, default);
        msg.ShouldNotBeNull();
        msg!.Subject.ShouldBe("foo");
        msg.Payload.ToArray().ShouldBe("hello"u8.ToArray());

        // The active block should have a write cache populated.
        // We verify this indirectly: after clearing, the read should still work (disk path).
        // BlockCount == 1 means there is exactly one block (the active one).
        store.BlockCount.ShouldBe(1);
    }

    // Go: TestFileStoreClearCache — filestore_test.go (clearCache eviction)
    [Fact]
    public async Task WriteCache_ClearEvictsButReadStillWorks()
    {
        // Write cache is an optimisation: clearing it should not affect correctness.
        // After clearing, reads fall through to disk and return the same data.
        await using var store = CreateStore(sub: "clear-cache");

        var seq = await store.AppendAsync("bar", "world"u8.ToArray(), default);

        // Access the single block directly via MsgBlock.Create/Recover round-trip:
        // We test ClearCache by writing several messages to force a block rotation
        // (the previous block's cache is cleared on rotation).

        // Write enough data to fill the first block and trigger rotation.
        var opts = new FileStoreOptions
        {
            Directory = Path.Combine(_dir, "rotate-test"),
            BlockSizeBytes = 256, // small block so rotation happens quickly
        };
        await using var storeSmall = CreateStore(opts);

        // Write several messages; block rotation will clear the cache on the sealed block.
        for (var i = 0; i < 10; i++)
            await storeSmall.AppendAsync($"sub.{i}", Encoding.UTF8.GetBytes($"payload-{i}"), default);

        // All messages should still be readable even though earlier blocks were sealed
        // and their caches were cleared.
        for (ulong s = 1; s <= 10; s++)
        {
            var m = await storeSmall.LoadAsync(s, default);
            m.ShouldNotBeNull();
        }
    }

    // -------------------------------------------------------------------------
    // TTL wheel tests
    // -------------------------------------------------------------------------

    // Go: TestFileStoreTtlWheelExpiry — filestore.go expireMsgs (thw.ExpireTasks)
    [Fact]
    public async Task TtlWheel_ExpiredMessagesRemoved()
    {
        // MaxAgeMs = 50ms: messages older than 50ms should be expired on the next append.
        var opts = new FileStoreOptions { MaxAgeMs = 50 };
        await using var store = CreateStore(opts, "ttl-expire");

        // Write some messages.
        await store.AppendAsync("events.a", "data-a"u8.ToArray(), default);
        await store.AppendAsync("events.b", "data-b"u8.ToArray(), default);

        var stateBefore = await store.GetStateAsync(default);
        stateBefore.Messages.ShouldBe(2UL);

        // Wait longer than the TTL.
        await Task.Delay(150);

        // Trigger expiry by appending a new message (expiry check happens at the start of each append).
        await store.AppendAsync("events.c", "data-c"u8.ToArray(), default);

        // The two old messages should now be gone; only the new one should remain.
        var stateAfter = await store.GetStateAsync(default);
        stateAfter.Messages.ShouldBe(1UL);
        stateAfter.LastSeq.ShouldBe(3UL);
    }

    // Go: TestFileStoreTtlWheelRetention — filestore.go expireMsgs (no expiry when fresh)
    [Fact]
    public async Task TtlWheel_UnexpiredMessagesRetained()
    {
        // MaxAgeMs = 5000ms: messages written just now should not be expired immediately.
        var opts = new FileStoreOptions { MaxAgeMs = 5000 };
        await using var store = CreateStore(opts, "ttl-retain");

        await store.AppendAsync("keep.a", "payload-a"u8.ToArray(), default);
        await store.AppendAsync("keep.b", "payload-b"u8.ToArray(), default);
        await store.AppendAsync("keep.c", "payload-c"u8.ToArray(), default);

        // Trigger the expiry check path via another append.
        await store.AppendAsync("keep.d", "payload-d"u8.ToArray(), default);

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe(4UL, "all four messages should still be present");
    }

    // -------------------------------------------------------------------------
    // StoreMsg sync method tests
    // -------------------------------------------------------------------------

    // Go: TestFileStoreStoreMsg — filestore.go storeMsg returns (seq, ts)
    [Fact]
    public async Task StoreMsg_ReturnsSequenceAndTimestamp()
    {
        await using var store = CreateStore(sub: "storemsg-basic");

        var beforeNs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        var (seq, ts) = store.StoreMsg("orders.new", null, "order-data"u8.ToArray(), 0L);
        var afterNs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;

        seq.ShouldBe(1UL);
        ts.ShouldBeGreaterThanOrEqualTo(beforeNs);
        ts.ShouldBeLessThanOrEqualTo(afterNs);

        // Verify the message is retrievable.
        var loaded = await store.LoadAsync(seq, default);
        loaded.ShouldNotBeNull();
        loaded!.Subject.ShouldBe("orders.new");
        loaded.Payload.ToArray().ShouldBe("order-data"u8.ToArray());
    }

    // Go: TestFileStoreStoreMsg — filestore.go storeMsg with headers
    [Fact]
    public async Task StoreMsg_WithHeaders_CombinesHeadersAndPayload()
    {
        await using var store = CreateStore(sub: "storemsg-headers");

        var hdr = "NATS/1.0\r\nX-Custom: value\r\n\r\n"u8.ToArray();
        var body = "message-body"u8.ToArray();
        var (seq, ts) = store.StoreMsg("events.all", hdr, body, 0L);

        seq.ShouldBe(1UL);
        ts.ShouldBeGreaterThan(0L);

        // The stored payload should be the combination of headers + body.
        var loaded = await store.LoadAsync(seq, default);
        loaded.ShouldNotBeNull();
        loaded!.Payload.Length.ShouldBe(hdr.Length + body.Length);
    }

    // Go: TestFileStoreStoreMsgPerMsgTtl — filestore.go per-message TTL override
    [Fact]
    public async Task StoreMsg_WithTtl_ExpiresAfterDelay()
    {
        // No stream-level TTL — only per-message TTL.
        await using var store = CreateStore(sub: "storemsg-ttl");

        // 80ms TTL in nanoseconds.
        const long ttlNs = 80_000_000L;

        var (seq, _) = store.StoreMsg("expire.me", null, "short-lived"u8.ToArray(), ttlNs);
        seq.ShouldBe(1UL);

        // Verify it's present immediately.
        var before = await store.GetStateAsync(default);
        before.Messages.ShouldBe(1UL);

        // Wait for expiry.
        await Task.Delay(200);

        // Trigger expiry by calling StoreMsg again (which calls ExpireFromWheel internally).
        store.StoreMsg("permanent", null, "stays"u8.ToArray(), 0L);

        // The TTL'd message should be gone; only the permanent one remains.
        var after = await store.GetStateAsync(default);
        after.Messages.ShouldBe(1UL);
        after.LastSeq.ShouldBe(2UL);
    }

    // Go: TestFileStoreStoreMsg — multiple sequential StoreMsgs increment sequence
    [Fact]
    public async Task StoreMsg_MultipleMessages_SequenceIncrements()
    {
        await using var store = CreateStore(sub: "storemsg-multi");

        for (var i = 1; i <= 5; i++)
        {
            var (seq, _) = store.StoreMsg($"topic.{i}", null, Encoding.UTF8.GetBytes($"msg-{i}"), 0L);
            seq.ShouldBe((ulong)i);
        }

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe(5UL);
        state.LastSeq.ShouldBe(5UL);
    }

    // -------------------------------------------------------------------------
    // Recovery re-registration test
    // -------------------------------------------------------------------------

    // Go: filestore.go recoverMsgs — TTL re-registration on restart
    [Fact]
    public async Task Recovery_ReregistersUnexpiredTtls()
    {
        // Write messages with a 5-second TTL (well beyond the test duration).
        // After recovering the store, the messages should still be present.
        var dir = Path.Combine(_dir, "ttl-recovery");
        var opts = new FileStoreOptions
        {
            Directory = dir,
            MaxAgeMs = 5000, // 5 second TTL
        };

        ulong seqA, seqB;

        // First open: write messages.
        {
            await using var store = new FileStore(opts);
            seqA = await store.AppendAsync("topic.a", "payload-a"u8.ToArray(), default);
            seqB = await store.AppendAsync("topic.b", "payload-b"u8.ToArray(), default);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe(2UL);
        } // FileStore disposed here.

        // Second open: recovery should re-register TTLs and messages should still be present.
        {
            await using var recovered = new FileStore(opts);

            var state = await recovered.GetStateAsync(default);
            state.Messages.ShouldBe(2UL, "unexpired messages should survive recovery");

            var msgA = await recovered.LoadAsync(seqA, default);
            msgA.ShouldNotBeNull();
            msgA!.Subject.ShouldBe("topic.a");

            var msgB = await recovered.LoadAsync(seqB, default);
            msgB.ShouldNotBeNull();
            msgB!.Subject.ShouldBe("topic.b");
        }
    }

    // Go: filestore.go recoverMsgs — expired messages removed on recovery
    [Fact]
    public async Task Recovery_ExpiredMessagesRemovedOnReopen()
    {
        // Write messages with a very short TTL, wait for them to expire, then
        // reopen the store. The expired messages should be pruned at startup.
        var dir = Path.Combine(_dir, "ttl-recovery-expired");
        var opts = new FileStoreOptions
        {
            Directory = dir,
            MaxAgeMs = 50, // 50ms TTL
        };

        // First open: write messages.
        {
            await using var store = new FileStore(opts);
            await store.AppendAsync("expiring.a", "data-a"u8.ToArray(), default);
            await store.AppendAsync("expiring.b", "data-b"u8.ToArray(), default);
        }

        // Wait for TTL to elapse.
        await Task.Delay(200);

        // Second open: expired messages should be pruned during RecoverBlocks -> PruneExpired.
        {
            await using var recovered = new FileStore(opts);

            var state = await recovered.GetStateAsync(default);
            state.Messages.ShouldBe(0UL, "expired messages should be removed on recovery");
        }
    }
}