CBDD/tests/CBDD.Tests/CompressionOverflowTests.cs

using System.IO.Compression;
using System.IO.MemoryMappedFiles;
using ZB.MOM.WW.CBDD.Core.Compression;
using ZB.MOM.WW.CBDD.Core.Storage;
using ZB.MOM.WW.CBDD.Shared;

namespace ZB.MOM.WW.CBDD.Tests;

public class CompressionOverflowTests
{
    /// <summary>
    /// Tests insert compressed document spanning overflow pages should round trip.
    /// </summary>
    [Fact]
    public void Insert_CompressedDocumentSpanningOverflowPages_ShouldRoundTrip()
    {
        var dbPath = NewDbPath();
        var options = new CompressionOptions
        {
            EnableCompression = true,
            MinSizeBytes = 64,
            MinSavingsPercent = 0,
            Codec = CompressionCodec.Deflate,
            Level = CompressionLevel.Fastest
        };

        try
        {
            using var db = new TestDbContext(dbPath, TinyPageConfig(), options);

            var payload = BuildPayload(300_000);
            var id = db.Users.Insert(new User { Name = payload, Age = 40 });
            db.SaveChanges();

            var found = db.Users.FindById(id);
            found.ShouldNotBeNull();
            found.Name.ShouldBe(payload);

            var counts = CountSlotModes(db.Storage);
            counts.CompressedOverflow.ShouldBeGreaterThanOrEqualTo(1);
            counts.OverflowPages.ShouldBeGreaterThanOrEqualTo(1);
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    /// <summary>
    /// Tests update should transition across compression thresholds.
    /// </summary>
    [Fact]
    public void Update_ShouldTransitionAcrossCompressionThresholds()
    {
        var dbPath = NewDbPath();
        var options = new CompressionOptions
        {
            EnableCompression = true,
            MinSizeBytes = 2048,
            MinSavingsPercent = 0,
            Codec = CompressionCodec.Brotli,
            Level = CompressionLevel.Fastest
        };

        try
        {
            using var db = new TestDbContext(dbPath, TinyPageConfig(), options);

            var user = new User { Name = "small", Age = 1 };
            var id = db.Users.Insert(user);
            db.SaveChanges();

            CountSlotModes(db.Storage).Compressed.ShouldBe(0);

            user.Name = BuildPayload(120_000);
            db.Users.Update(user).ShouldBeTrue();
            db.SaveChanges();

            var afterLarge = db.Users.FindById(id);
            afterLarge.ShouldNotBeNull();
            afterLarge.Name.ShouldBe(user.Name);

            var largeCounts = CountSlotModes(db.Storage);
            largeCounts.Compressed.ShouldBeGreaterThanOrEqualTo(1);

            user.Name = "small-again";
            db.Users.Update(user).ShouldBeTrue();
            db.SaveChanges();

            var afterShrink = db.Users.FindById(id);
            afterShrink.ShouldNotBeNull();
            afterShrink.Name.ShouldBe("small-again");

            var finalCounts = CountSlotModes(db.Storage);
            finalCounts.Compressed.ShouldBe(0);
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    private static (int Compressed, int CompressedOverflow, int OverflowPages) CountSlotModes(StorageEngine storage)
    {
        var buffer = new byte[storage.PageSize];
        var compressed = 0;
        var compressedOverflow = 0;
        var overflowPages = 0;

        for (uint pageId = 1; pageId < storage.PageCount; pageId++)
        {
            storage.ReadPage(pageId, null, buffer);
            var header = SlottedPageHeader.ReadFrom(buffer);

            if (header.PageType == PageType.Overflow)
            {
                overflowPages++;
                continue;
            }

            if (header.PageType != PageType.Data)
                continue;

            for (ushort slotIndex = 0; slotIndex < header.SlotCount; slotIndex++)
            {
                var slotOffset = SlottedPageHeader.Size + (slotIndex * SlotEntry.Size);
                var slot = SlotEntry.ReadFrom(buffer.AsSpan(slotOffset, SlotEntry.Size));
                if ((slot.Flags & SlotFlags.Deleted) != 0)
                    continue;

                var isCompressed = (slot.Flags & SlotFlags.Compressed) != 0;
                var hasOverflow = (slot.Flags & SlotFlags.HasOverflow) != 0;
                if (isCompressed)
                    compressed++;
                if (isCompressed && hasOverflow)
                    compressedOverflow++;
            }
        }

        return (compressed, compressedOverflow, overflowPages);
    }

    private static PageFileConfig TinyPageConfig()
    {
        return new PageFileConfig
        {
            PageSize = 16 * 1024,
            InitialFileSize = 1024 * 1024,
            Access = MemoryMappedFileAccess.ReadWrite
        };
    }

    private static string BuildPayload(int approxLength)
    {
        var builder = new System.Text.StringBuilder(approxLength + 256);
        var i = 0;
        while (builder.Length < approxLength)
        {
            builder.Append("overflow-payload-");
            builder.Append(i.ToString("D7"));
            builder.Append('|');
            i++;
        }

        return builder.ToString();
    }

    private static string NewDbPath()
        => Path.Combine(Path.GetTempPath(), $"compression_overflow_{Guid.NewGuid():N}.db");

    private static void CleanupFiles(string dbPath)
    {
        var walPath = Path.ChangeExtension(dbPath, ".wal");
        var markerPath = $"{dbPath}.compact.state";
        if (File.Exists(dbPath)) File.Delete(dbPath);
        if (File.Exists(walPath)) File.Delete(walPath);
        if (File.Exists(markerPath)) File.Delete(markerPath);
    }
}