CBDD/tests/CBDD.Tests/CompactionCrashRecoveryTests.cs

using ZB.MOM.WW.CBDD.Bson;
using ZB.MOM.WW.CBDD.Core.Storage;
using ZB.MOM.WW.CBDD.Shared;

namespace ZB.MOM.WW.CBDD.Tests;

public class CompactionCrashRecoveryTests
{
    [Theory]
    [InlineData("Started")]
    [InlineData("Copied")]
    [InlineData("Swapped")]
    public void ResumeCompaction_FromCrashMarkerPhases_ShouldFinalizeAndPreserveData(string phase)
    {
        var dbPath = NewDbPath();
        var markerPath = MarkerPath(dbPath);

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

            var ids = SeedData(db);
            db.ForceCheckpoint();
            WriteMarker(markerPath, dbPath, phase);

            var resumed = db.Storage.ResumeCompactionIfNeeded(new CompactionOptions
            {
                EnableTailTruncation = true,
                DefragmentSlottedPages = true,
                NormalizeFreeList = true
            });

            resumed.ShouldNotBeNull();
            resumed!.ResumedFromMarker.ShouldBeTrue();
            File.Exists(markerPath).ShouldBeFalse();

            db.Users.Count().ShouldBe(ids.Count);
            var recoveredDoc = ids
                .Select(id => db.Users.FindById(id))
                .FirstOrDefault(x => x != null);
            recoveredDoc.ShouldNotBeNull();
            recoveredDoc!.Name.ShouldContain("user-");

            db.Storage.ResumeCompactionIfNeeded().ShouldBeNull();
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    [Fact]
    public void ResumeCompaction_WithCorruptedMarker_ShouldRecoverDeterministically()
    {
        var dbPath = NewDbPath();
        var markerPath = MarkerPath(dbPath);

        try
        {
            using var db = new TestDbContext(dbPath);
            var ids = SeedData(db);
            db.ForceCheckpoint();

            File.WriteAllText(markerPath, "{invalid-json-marker");

            var resumed = db.Storage.ResumeCompactionIfNeeded(new CompactionOptions
            {
                EnableTailTruncation = true
            });

            resumed.ShouldNotBeNull();
            resumed!.ResumedFromMarker.ShouldBeTrue();
            File.Exists(markerPath).ShouldBeFalse();

            db.Users.Count().ShouldBe(ids.Count);
            var recoveredDoc = ids
                .Select(id => db.Users.FindById(id))
                .FirstOrDefault(x => x != null);
            recoveredDoc.ShouldNotBeNull();
            recoveredDoc!.Name.ShouldContain("user-");
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    private static List<ObjectId> SeedData(TestDbContext db)
    {
        var ids = new List<ObjectId>();
        for (var i = 0; i < 120; i++)
        {
            ids.Add(db.Users.Insert(new User
            {
                Name = $"user-{i:D4}-payload-{new string('x', 120)}",
                Age = i % 20
            }));
        }

        db.SaveChanges();
        return ids;
    }

    private static void WriteMarker(string markerPath, string dbPath, string phase)
    {
        var safeDbPath = dbPath.Replace("\\", "\\\\", StringComparison.Ordinal);
        var now = DateTimeOffset.UtcNow.ToString("O");
        var json = $$"""
            {"version":1,"phase":"{{phase}}","databasePath":"{{safeDbPath}}","startedAtUtc":"{{now}}","lastUpdatedUtc":"{{now}}","onlineMode":false,"mode":"InPlace"}
            """;
        File.WriteAllText(markerPath, json);
    }

    private static string MarkerPath(string dbPath) => $"{dbPath}.compact.state";

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

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