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

// Reference: golang/nats-server/server/filestore_test.go
// Tests ported from: TestFileStoreRemovePartialRecovery,
//                    TestFileStoreRemoveOutOfOrderRecovery,
//                    TestFileStoreAgeLimitRecovery, TestFileStoreBitRot,
//                    TestFileStoreEraseAndNoIndexRecovery,
//                    TestFileStoreExpireMsgsOnStart,
//                    TestFileStoreRebuildStateDmapAccountingBug,
//                    TestFileStoreRecalcFirstSequenceBug,
//                    TestFileStoreFullStateBasics

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

namespace NATS.Server.Tests.JetStream.Storage;

public sealed class FileStoreRecoveryTests : IDisposable
{
    private readonly string _dir;

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

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

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

    // Go: TestFileStoreRemovePartialRecovery server/filestore_test.go:1076
    [Fact]
    public async Task Remove_half_then_recover()
    {
        var subDir = "partial-recovery";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 100; i++)
                await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);

            // Remove first half.
            for (ulong i = 1; i <= 50; i++)
                await store.RemoveAsync(i, default);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)50);
        }

        // Recover and verify state matches.
        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)50);
            state.FirstSeq.ShouldBe((ulong)51);
            state.LastSeq.ShouldBe((ulong)100);

            // Verify removed messages are gone.
            for (ulong i = 1; i <= 50; i++)
                (await store.LoadAsync(i, default)).ShouldBeNull();

            // Verify remaining messages are present.
            for (ulong i = 51; i <= 100; i++)
                (await store.LoadAsync(i, default)).ShouldNotBeNull();
        }
    }

    // Go: TestFileStoreRemoveOutOfOrderRecovery server/filestore_test.go:1119
    [Fact]
    public async Task Remove_evens_then_recover()
    {
        var subDir = "ooo-recovery";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 100; i++)
                await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);

            // Remove even-numbered sequences.
            for (var i = 2; i <= 100; i += 2)
                (await store.RemoveAsync((ulong)i, default)).ShouldBeTrue();

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)50);
        }

        // Recover and verify.
        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)50);

            // Seq 1 should exist.
            (await store.LoadAsync(1, default)).ShouldNotBeNull();

            // Even sequences should be gone.
            for (var i = 2; i <= 100; i += 2)
                (await store.LoadAsync((ulong)i, default)).ShouldBeNull();

            // Odd sequences should exist.
            for (var i = 1; i <= 99; i += 2)
                (await store.LoadAsync((ulong)i, default)).ShouldNotBeNull();
        }
    }

    // Go: TestFileStoreAgeLimitRecovery server/filestore_test.go:1183
    [Fact]
    public async Task Age_limit_recovery_expires_on_restart()
    {
        var subDir = "age-recovery";

        await using (var store = CreateStore(subDir, new FileStoreOptions { MaxAgeMs = 200 }))
        {
            for (var i = 0; i < 20; i++)
                await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);

            (await store.GetStateAsync(default)).Messages.ShouldBe((ulong)20);
        }

        // Wait for messages to age out.
        await Task.Delay(300);

        // Reopen — expired messages should be pruned on load.
        await using (var store = CreateStore(subDir, new FileStoreOptions { MaxAgeMs = 200 }))
        {
            // Trigger prune by appending.
            await store.AppendAsync("foo", "trigger"u8.ToArray(), default);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)1);
        }
    }

    // Go: TestFileStoreEraseAndNoIndexRecovery server/filestore_test.go:1363
    [Fact]
    public async Task Remove_evens_then_recover_without_index()
    {
        var subDir = "no-index-recovery";
        var dir = Path.Combine(_dir, subDir);

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 100; i++)
                await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);

            // Remove even-numbered sequences.
            for (var i = 2; i <= 100; i += 2)
                (await store.RemoveAsync((ulong)i, default)).ShouldBeTrue();

            (await store.GetStateAsync(default)).Messages.ShouldBe((ulong)50);
        }

        // Remove the index manifest file to force a full rebuild.
        var manifestPath = Path.Combine(dir, "index.manifest.json");
        if (File.Exists(manifestPath))
            File.Delete(manifestPath);

        // Recover without index manifest.
        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)50);

            // Even sequences should still be gone.
            for (var i = 2; i <= 100; i += 2)
                (await store.LoadAsync((ulong)i, default)).ShouldBeNull();

            // Odd sequences should exist.
            for (var i = 1; i <= 99; i += 2)
                (await store.LoadAsync((ulong)i, default)).ShouldNotBeNull();
        }
    }

    // Go: TestFileStoreBitRot server/filestore_test.go:1229
    [Fact]
    public async Task Corrupted_data_file_loses_messages_but_store_recovers()
    {
        var subDir = "bitrot";
        var dir = Path.Combine(_dir, subDir);

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 20; i++)
                await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);
        }

        // Corrupt the data file by writing random bytes in the middle.
        var dataFile = Path.Combine(dir, "messages.jsonl");
        if (File.Exists(dataFile))
        {
            var content = File.ReadAllBytes(dataFile);
            if (content.Length > 50)
            {
                // Corrupt some bytes in the middle.
                content[content.Length / 2] = 0xFF;
                content[content.Length / 2 + 1] = 0xFE;
                File.WriteAllBytes(dataFile, content);
            }
        }

        // Recovery should not throw; it may lose some messages though.
        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            // We may lose messages due to corruption, but at least some should survive
            // if the corruption only affected one record.
            // The key point is that the store recovered without throwing.
            state.Messages.ShouldBeGreaterThanOrEqualTo((ulong)0);
        }
    }

    // Go: TestFileStoreFullStateBasics server/filestore_test.go:5461
    [Fact]
    public async Task Full_state_recovery_preserves_all_messages()
    {
        var subDir = "full-state";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 50; i++)
                await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i}"), default);

            for (var i = 0; i < 50; i++)
                await store.AppendAsync("bar", Encoding.UTF8.GetBytes($"msg-{i}"), default);
        }

        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)100);
            state.FirstSeq.ShouldBe((ulong)1);
            state.LastSeq.ShouldBe((ulong)100);

            var msg1 = await store.LoadAsync(1, default);
            msg1.ShouldNotBeNull();
            msg1!.Subject.ShouldBe("foo");

            var msg51 = await store.LoadAsync(51, default);
            msg51.ShouldNotBeNull();
            msg51!.Subject.ShouldBe("bar");
        }
    }

    // Go: TestFileStoreExpireMsgsOnStart server/filestore_test.go:3018
    [Fact]
    public async Task Expire_on_restart_with_different_maxage()
    {
        var subDir = "expire-on-start";

        // Store with no age limit.
        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 10; i++)
                await store.AppendAsync("foo", "Hello"u8.ToArray(), default);
        }

        await Task.Delay(100);

        // Reopen with an age limit that will expire all old messages.
        await using (var store = CreateStore(subDir, new FileStoreOptions { MaxAgeMs = 50 }))
        {
            // Trigger pruning.
            await store.AppendAsync("foo", "trigger"u8.ToArray(), default);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)1);
        }
    }

    // Go: TestFileStoreRemovePartialRecovery server/filestore_test.go:1076
    [Fact]
    public async Task Remove_then_append_then_recover()
    {
        var subDir = "rm-append-recover";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 10; i++)
                await store.AppendAsync("foo", "Hello"u8.ToArray(), default);

            await store.RemoveAsync(5, default);
            await store.AppendAsync("foo", "After remove"u8.ToArray(), default);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)10);
            state.LastSeq.ShouldBe((ulong)11);
        }

        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)10);
            state.LastSeq.ShouldBe((ulong)11);

            (await store.LoadAsync(5, default)).ShouldBeNull();
            (await store.LoadAsync(11, default)).ShouldNotBeNull();
        }
    }

    // Go: TestFileStoreRecalcFirstSequenceBug server/filestore_test.go:5405
    [Fact]
    public async Task Recovery_preserves_first_seq_after_removes()
    {
        var subDir = "first-seq-recovery";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 20; i++)
                await store.AppendAsync("foo", "data"u8.ToArray(), default);

            // Remove first 10.
            for (ulong i = 1; i <= 10; i++)
                await store.RemoveAsync(i, default);

            var state = await store.GetStateAsync(default);
            state.FirstSeq.ShouldBe((ulong)11);
        }

        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.FirstSeq.ShouldBe((ulong)11);
            state.Messages.ShouldBe((ulong)10);
        }
    }

    // Go: TestFileStoreRebuildStateDmapAccountingBug server/filestore_test.go:3692
    [Fact]
    public async Task Recovery_with_scattered_deletes_preserves_count()
    {
        var subDir = "scattered-deletes";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 50; i++)
                await store.AppendAsync("foo", "data"u8.ToArray(), default);

            // Delete scattered: every 3rd.
            for (var i = 3; i <= 50; i += 3)
                await store.RemoveAsync((ulong)i, default);

            var expectedCount = 50 - (50 / 3);

            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)expectedCount);
        }

        await using (var store = CreateStore(subDir))
        {
            var expectedCount = 50 - (50 / 3);
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)expectedCount);
        }
    }

    // Go: TestFileStoreBasicWriteMsgsAndRestore server/filestore_test.go:181
    [Fact]
    public async Task Recovery_preserves_message_payloads()
    {
        var subDir = "payload-recovery";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 10; i++)
                await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"message-{i}"), default);
        }

        await using (var store = CreateStore(subDir))
        {
            for (ulong i = 1; i <= 10; i++)
            {
                var msg = await store.LoadAsync(i, default);
                msg.ShouldNotBeNull();
                msg!.Subject.ShouldBe("foo");
                var expected = Encoding.UTF8.GetBytes($"message-{i - 1}");
                msg.Payload.ToArray().ShouldBe(expected);
            }
        }
    }

    // Go: TestFileStoreBasicWriteMsgsAndRestore server/filestore_test.go:181
    [Fact]
    public async Task Recovery_preserves_subjects()
    {
        var subDir = "subject-recovery";

        await using (var store = CreateStore(subDir))
        {
            await store.AppendAsync("alpha", "one"u8.ToArray(), default);
            await store.AppendAsync("beta", "two"u8.ToArray(), default);
            await store.AppendAsync("gamma", "three"u8.ToArray(), default);
        }

        await using (var store = CreateStore(subDir))
        {
            var msg1 = await store.LoadAsync(1, default);
            msg1.ShouldNotBeNull();
            msg1!.Subject.ShouldBe("alpha");

            var msg2 = await store.LoadAsync(2, default);
            msg2.ShouldNotBeNull();
            msg2!.Subject.ShouldBe("beta");

            var msg3 = await store.LoadAsync(3, default);
            msg3.ShouldNotBeNull();
            msg3!.Subject.ShouldBe("gamma");
        }
    }

    // Go: TestFileStoreRemoveOutOfOrderRecovery server/filestore_test.go:1119
    [Fact]
    public async Task Recovery_with_large_message_count()
    {
        var subDir = "large-recovery";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 500; i++)
                await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i:D4}"), default);
        }

        await using (var store = CreateStore(subDir))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)500);
            state.FirstSeq.ShouldBe((ulong)1);
            state.LastSeq.ShouldBe((ulong)500);
        }
    }
}