CBDD/tests/CBDD.Tests/CompressionCompatibilityTests.cs

using System.IO.Compression;
using System.Security.Cryptography;
using ZB.MOM.WW.CBDD.Bson;
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 CompressionCompatibilityTests
{
    [Fact]
    public void OpeningLegacyUncompressedFile_WithCompressionEnabled_ShouldNotMutateDbFile()
    {
        var dbPath = NewDbPath();
        var idList = new List<ObjectId>();

        try
        {
            using (var db = new TestDbContext(dbPath))
            {
                idList.Add(db.Users.Insert(new User { Name = "legacy-a", Age = 10 }));
                idList.Add(db.Users.Insert(new User { Name = "legacy-b", Age = 11 }));
                db.SaveChanges();
                db.ForceCheckpoint();
            }

            var beforeSize = new FileInfo(dbPath).Length;
            var beforeHash = ComputeFileHash(dbPath);

            var compressionOptions = new CompressionOptions
            {
                EnableCompression = true,
                MinSizeBytes = 0,
                MinSavingsPercent = 0,
                Codec = CompressionCodec.Brotli,
                Level = CompressionLevel.Fastest
            };

            using (var reopened = new TestDbContext(dbPath, compressionOptions))
            {
                reopened.Users.FindById(idList[0])!.Name.ShouldBe("legacy-a");
                reopened.Users.FindById(idList[1])!.Name.ShouldBe("legacy-b");
                reopened.Users.Count().ShouldBe(2);
            }

            var afterSize = new FileInfo(dbPath).Length;
            var afterHash = ComputeFileHash(dbPath);

            afterSize.ShouldBe(beforeSize);
            afterHash.ShouldBe(beforeHash);
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    [Fact]
    public void MixedFormatDocuments_ShouldRemainReadableAfterPartialMigration()
    {
        var dbPath = NewDbPath();
        ObjectId legacyId;
        ObjectId compressedId;

        try
        {
            using (var db = new TestDbContext(dbPath))
            {
                legacyId = db.Users.Insert(new User { Name = "legacy-uncompressed", Age = 22 });
                db.SaveChanges();
                db.ForceCheckpoint();
            }

            var compressionOptions = new CompressionOptions
            {
                EnableCompression = true,
                MinSizeBytes = 0,
                MinSavingsPercent = 0,
                Codec = CompressionCodec.Brotli,
                Level = CompressionLevel.Fastest
            };

            using (var migrated = new TestDbContext(dbPath, compressionOptions))
            {
                compressedId = migrated.Users.Insert(new User { Name = BuildPayload(24_000), Age = 33 });
                migrated.SaveChanges();
                migrated.ForceCheckpoint();
            }

            using (var verify = new TestDbContext(dbPath, compressionOptions))
            {
                verify.Users.FindById(legacyId)!.Name.ShouldBe("legacy-uncompressed");
                verify.Users.FindById(compressedId)!.Name.Length.ShouldBeGreaterThan(10_000);

                var counts = CountActiveDataSlots(verify.Storage);
                counts.Compressed.ShouldBeGreaterThanOrEqualTo(1);
                counts.Uncompressed.ShouldBeGreaterThanOrEqualTo(1);
            }
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    private static (int Compressed, int Uncompressed) CountActiveDataSlots(StorageEngine storage)
    {
        var buffer = new byte[storage.PageSize];
        var compressed = 0;
        var uncompressed = 0;

        for (uint pageId = 1; pageId < storage.PageCount; pageId++)
        {
            storage.ReadPage(pageId, null, buffer);
            var header = SlottedPageHeader.ReadFrom(buffer);
            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;

                if ((slot.Flags & SlotFlags.Compressed) != 0)
                    compressed++;
                else
                    uncompressed++;
            }
        }

        return (compressed, uncompressed);
    }

    private static string ComputeFileHash(string path)
    {
        using var stream = File.OpenRead(path);
        using var sha256 = SHA256.Create();
        return Convert.ToHexString(sha256.ComputeHash(stream));
    }

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

        return builder.ToString();
    }

    private static string NewDbPath()
        => Path.Combine(Path.GetTempPath(), $"compression_compat_{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);
    }
}