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

// Reference: golang/nats-server/server/filestore_test.go
// Tests ported from: TestFileStoreBasics (S2Compression permutation),
//                    TestFileStoreWriteExpireWrite (compression variant),
//                    TestFileStoreAgeLimit (compression variant),
//                    TestFileStoreCompactLastPlusOne (compression variant)
// The Go tests use testFileStoreAllPermutations to run each test with
// NoCompression and S2Compression. These tests exercise the .NET compression path.

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

namespace NATS.Server.Tests.JetStream.Storage;

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

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

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

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

    // Go: TestFileStoreBasics server/filestore_test.go:86 (S2 permutation)
    [Fact]
    public async Task Compressed_store_and_load()
    {
        await using var store = CreateStore("comp-basic");

        const string subject = "foo";
        var payload = "Hello World"u8.ToArray();

        for (var i = 1; i <= 5; i++)
        {
            var seq = await store.AppendAsync(subject, payload, default);
            seq.ShouldBe((ulong)i);
        }

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

        var msg = await store.LoadAsync(3, default);
        msg.ShouldNotBeNull();
        msg!.Subject.ShouldBe(subject);
        msg.Payload.ToArray().ShouldBe(payload);
    }

    // Go: TestFileStoreBasicWriteMsgsAndRestore server/filestore_test.go:181 (S2 permutation)
    [Fact]
    public async Task Compressed_store_and_recover()
    {
        var subDir = "comp-recover";

        await using (var store = CreateStore(subDir))
        {
            for (var i = 0; i < 100; 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)100);

            var msg = await store.LoadAsync(50, default);
            msg.ShouldNotBeNull();
            msg!.Subject.ShouldBe("foo");
            msg.Payload.ToArray().ShouldBe(Encoding.UTF8.GetBytes("msg-0049"));
        }
    }

    // Go: TestFileStoreBasics server/filestore_test.go:86 (S2 permutation)
    [Fact]
    public async Task Compressed_remove_and_reload()
    {
        await using var store = CreateStore("comp-remove");

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

        await store.RemoveAsync(5, default);

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

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

    // Go: TestFileStorePurge server/filestore_test.go:709 (S2 permutation)
    [Fact]
    public async Task Compressed_purge()
    {
        await using var store = CreateStore("comp-purge");

        for (var i = 0; i < 20; i++)
            await store.AppendAsync("foo", "Hello"u8.ToArray(), default);

        await store.PurgeAsync(default);

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

    // Go: TestFileStoreWriteExpireWrite server/filestore_test.go:424 (S2 permutation)
    [Fact]
    public async Task Compressed_large_batch()
    {
        await using var store = CreateStore("comp-large");

        for (var i = 0; i < 200; i++)
            await store.AppendAsync("zzz", Encoding.UTF8.GetBytes($"Hello World! - {i}"), default);

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

        for (ulong i = 1; i <= 200; i++)
        {
            var msg = await store.LoadAsync(i, default);
            msg.ShouldNotBeNull();
        }
    }

    // Go: TestFileStoreAgeLimit server/filestore_test.go:616 (S2 permutation)
    [Fact]
    public async Task Compressed_with_age_expiry()
    {
        await using var store = CreateStore("comp-age", options: new FileStoreOptions { MaxAgeMs = 200 });

        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", "Hello"u8.ToArray(), default);

        await Task.Delay(300);

        await store.AppendAsync("foo", "trigger"u8.ToArray(), default);

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

    // Go: TestFileStoreSnapshot server/filestore_test.go:1799 (S2 permutation)
    [Fact]
    public async Task Compressed_snapshot_and_restore()
    {
        await using var store = CreateStore("comp-snap-src");

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

        var snap = await store.CreateSnapshotAsync(default);
        snap.Length.ShouldBeGreaterThan(0);

        await using var restored = CreateStore("comp-snap-dst");
        await restored.RestoreSnapshotAsync(snap, default);

        var srcState = await store.GetStateAsync(default);
        var dstState = await restored.GetStateAsync(default);
        dstState.Messages.ShouldBe(srcState.Messages);

        for (ulong i = 1; i <= srcState.Messages; i++)
        {
            var original = await store.LoadAsync(i, default);
            var copy = await restored.LoadAsync(i, default);
            copy.ShouldNotBeNull();
            copy!.Payload.ToArray().ShouldBe(original!.Payload.ToArray());
        }
    }

    // Combined encryption + compression (Go AES-S2 permutation).
    [Fact]
    public async Task Compressed_and_encrypted_round_trip()
    {
        var dir = Path.Combine(_dir, "comp-enc");
        await using var store = new FileStore(new FileStoreOptions
        {
            Directory = dir,
            EnableCompression = true,
            EnableEncryption = true,
            EncryptionKey = "test-key-for-compression!!!!!!"u8.ToArray(),
        });

        var payload = "Hello World - compressed and encrypted"u8.ToArray();
        for (var i = 0; i < 10; i++)
            await store.AppendAsync("foo", payload, default);

        for (ulong i = 1; i <= 10; i++)
        {
            var msg = await store.LoadAsync(i, default);
            msg.ShouldNotBeNull();
            msg!.Payload.ToArray().ShouldBe(payload);
        }
    }

    // Combined encryption + compression with recovery.
    [Fact]
    public async Task Compressed_and_encrypted_recovery()
    {
        var subDir = "comp-enc-recover";
        var dir = Path.Combine(_dir, subDir);
        var key = "test-key-for-compression!!!!!!"u8.ToArray();

        await using (var store = new FileStore(new FileStoreOptions
        {
            Directory = dir,
            EnableCompression = true,
            EnableEncryption = true,
            EncryptionKey = key,
        }))
        {
            for (var i = 0; i < 20; i++)
                await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i:D4}"), default);
        }

        await using (var store = new FileStore(new FileStoreOptions
        {
            Directory = dir,
            EnableCompression = true,
            EnableEncryption = true,
            EncryptionKey = key,
        }))
        {
            var state = await store.GetStateAsync(default);
            state.Messages.ShouldBe((ulong)20);

            var msg = await store.LoadAsync(15, default);
            msg.ShouldNotBeNull();
            msg!.Payload.ToArray().ShouldBe(Encoding.UTF8.GetBytes("msg-0014"));
        }
    }

    // Compressed large payload (highly compressible).
    [Fact]
    public async Task Compressed_highly_compressible_payload()
    {
        await using var store = CreateStore("comp-compressible");

        // Highly repetitive data should compress well.
        var payload = new byte[4096];
        Array.Fill(payload, (byte)'A');

        await store.AppendAsync("foo", payload, default);

        var msg = await store.LoadAsync(1, default);
        msg.ShouldNotBeNull();
        msg!.Payload.ToArray().ShouldBe(payload);
    }

    // Compressed empty payload.
    [Fact]
    public async Task Compressed_empty_payload()
    {
        await using var store = CreateStore("comp-empty");

        await store.AppendAsync("foo", ReadOnlyMemory<byte>.Empty, default);

        var msg = await store.LoadAsync(1, default);
        msg.ShouldNotBeNull();
        msg!.Payload.Length.ShouldBe(0);
    }

    // Verify compressed data is different from uncompressed on disk.
    [Fact]
    public async Task Compressed_data_differs_from_uncompressed_on_disk()
    {
        var compDir = Path.Combine(_dir, "comp-on-disk");
        var plainDir = Path.Combine(_dir, "plain-on-disk");

        await using (var compStore = CreateStore("comp-on-disk"))
        {
            await compStore.AppendAsync("foo", "AAAAAAAAAAAAAAAAAAAAAAAAAAA"u8.ToArray(), default);
        }

        await using (var plainStore = CreateStore("plain-on-disk", compress: false))
        {
            await plainStore.AppendAsync("foo", "AAAAAAAAAAAAAAAAAAAAAAAAAAA"u8.ToArray(), default);
        }

        var compFile = Path.Combine(compDir, "messages.jsonl");
        var plainFile = Path.Combine(plainDir, "messages.jsonl");

        if (File.Exists(compFile) && File.Exists(plainFile))
        {
            var compContent = File.ReadAllText(compFile);
            var plainContent = File.ReadAllText(plainFile);
            // The base64-encoded payloads should differ due to compression envelope.
            compContent.ShouldNotBe(plainContent);
        }
    }
}