smartbotic-database/tests/test_snapshot_durability.cpp

/**
 * Integration tests for v1.6.1 snapshot durability + fallback.
 *
 * Exercises SnapshotManager directly (no gRPC server).
 *
 * Test cases:
 *   1. atomic write leaves no .tmp file behind
 *   2. snapshot round-trip (write then load back)
 *   3. post-write verification catches externally-truncated file
 *   4. loadWithFallback uses older snapshot when newest is corrupt
 *   5. orphaned .tmp files cleaned by listSnapshots
 *   6. loadWithFallback throws when all snapshots corrupt
 */

#include "../service/src/persistence/snapshot.hpp"
#include "../service/src/persistence/wal.hpp"
#include "../service/src/memory_store.hpp"
#include "../service/src/document.hpp"

#include <cassert>
#include <chrono>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <nlohmann/json.hpp>
#include <thread>
#include <unistd.h>  // getpid
#include <unordered_set>

namespace fs = std::filesystem;
using smartbotic::database::MemoryStore;
using smartbotic::database::SnapshotManager;
using smartbotic::database::Document;
using smartbotic::database::CollectionOptions;

namespace {

MemoryStore::Config defaultStoreConfig() {
    MemoryStore::Config cfg;
    cfg.nodeId = "test";
    cfg.maxMemoryBytes = 256ULL * 1024 * 1024;
    return cfg;
}

fs::path makeSnapshotDir(const std::string& name) {
    auto dir = fs::temp_directory_path() /
               ("snap-durability-" + name + "-" + std::to_string(::getpid()));
    fs::remove_all(dir);
    fs::create_directories(dir);
    return dir;
}

Document makeDoc(const std::string& id) {
    Document d;
    d.id = id;
    d.set_data(nlohmann::json{{"value", id}});
    return d;
}

} // anonymous namespace

// ──────────────────────────────────────────────────────────
// Test 1: atomic write leaves no .tmp file behind
// ──────────────────────────────────────────────────────────
void test_atomic_write_leaves_no_tmp() {
    auto dir = makeSnapshotDir("atomic");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    scfg.validateAfterWrite = true;
    SnapshotManager mgr(scfg);

    MemoryStore store(defaultStoreConfig());
    store.start();
    store.createCollection("docs", CollectionOptions{});
    for (int i = 0; i < 50; i++) {
        store.insert("docs", makeDoc("d" + std::to_string(i)));
    }

    auto path = mgr.createSnapshot(store, 42);

    // No .tmp files should remain
    int tmpCount = 0;
    for (auto& e : fs::directory_iterator(dir)) {
        if (e.path().extension() == ".tmp") tmpCount++;
    }
    assert(tmpCount == 0);

    // The final file should exist and pass verification
    assert(fs::exists(path));
    assert(mgr.verifySnapshot(path));

    store.stop();
    fs::remove_all(dir);
    std::cout << "PASS: atomic write leaves no .tmp file behind\n";
}

// ──────────────────────────────────────────────────────────
// Test 2: snapshot round-trips (write then load back)
// ──────────────────────────────────────────────────────────
void test_snapshot_round_trip() {
    auto dir = makeSnapshotDir("roundtrip");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    SnapshotManager mgr(scfg);

    MemoryStore store(defaultStoreConfig());
    store.start();
    store.createCollection("docs", CollectionOptions{});
    for (int i = 0; i < 20; i++) {
        store.insert("docs", makeDoc("d" + std::to_string(i)));
    }
    auto path = mgr.createSnapshot(store, 99);
    store.stop();

    MemoryStore store2(defaultStoreConfig());
    store2.start();
    uint64_t seq = mgr.loadSnapshot(path, store2);
    assert(seq == 99);
    auto doc = store2.get("docs", "d5");
    assert(doc.has_value());
    store2.stop();

    fs::remove_all(dir);
    std::cout << "PASS: snapshot round-trip\n";
}

// ──────────────────────────────────────────────────────────
// Test 3: post-write verification catches externally-truncated file
// ──────────────────────────────────────────────────────────
void test_verification_catches_truncation() {
    auto dir = makeSnapshotDir("truncate");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    SnapshotManager mgr(scfg);

    MemoryStore store(defaultStoreConfig());
    store.start();
    store.createCollection("docs", CollectionOptions{});
    for (int i = 0; i < 50; i++) {
        store.insert("docs", makeDoc("d" + std::to_string(i)));
    }
    auto path = mgr.createSnapshot(store, 1);
    store.stop();

    // Truncate the file AFTER a successful write to simulate the Zoe scenario
    uintmax_t original = fs::file_size(path);
    fs::resize_file(path, original / 2);

    // verifySnapshot should return false
    assert(!mgr.verifySnapshot(path));

    fs::remove_all(dir);
    std::cout << "PASS: post-write verification catches truncation\n";
}

// ──────────────────────────────────────────────────────────
// Test 4: loadWithFallback uses older snapshot when newest is corrupt
// ──────────────────────────────────────────────────────────
void test_loadWithFallback_uses_older_when_newest_corrupt() {
    auto dir = makeSnapshotDir("fallback");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    scfg.maxSnapshots = 10;
    SnapshotManager mgr(scfg);

    MemoryStore store(defaultStoreConfig());
    store.start();
    store.createCollection("docs", CollectionOptions{});

    // Older snapshot — 10 docs
    for (int i = 0; i < 10; i++) {
        store.insert("docs", makeDoc("d" + std::to_string(i)));
    }
    auto older = mgr.createSnapshot(store, 100);

    // Filenames include timestamps down to the second, so sleep to ensure ordering
    std::this_thread::sleep_for(std::chrono::seconds(1));

    // Newer snapshot — 20 docs
    for (int i = 10; i < 20; i++) {
        store.insert("docs", makeDoc("d" + std::to_string(i)));
    }
    auto newer = mgr.createSnapshot(store, 200);
    store.stop();

    assert(older != newer);

    // Corrupt the newer
    fs::resize_file(newer, fs::file_size(newer) / 2);
    assert(!mgr.verifySnapshot(newer));

    // loadWithFallback should transparently fall back to the older
    MemoryStore store2(defaultStoreConfig());
    store2.start();
    fs::path used;
    size_t attempted = 0;
    uint64_t seq = mgr.loadWithFallback(store2, &used, &attempted);
    assert(seq == 100);
    assert(used == older);
    assert(attempted == 2);  // tried newer first, then older
    store2.stop();

    fs::remove_all(dir);
    std::cout << "PASS: loadWithFallback uses older snapshot when newest corrupt\n";
}

// ──────────────────────────────────────────────────────────
// Test 5: orphaned .tmp files cleaned by listSnapshots()
// ──────────────────────────────────────────────────────────
void test_orphaned_tmp_cleaned_by_listSnapshots() {
    auto dir = makeSnapshotDir("orphan");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    SnapshotManager mgr(scfg);

    // Manually plant a fake .tmp like a crashed-mid-write would leave
    auto tmp = dir / "snapshot-20260101-120000.dat.tmp";
    {
        std::ofstream f(tmp);
        f << "partial";
    }
    assert(fs::exists(tmp));

    // listSnapshots should clean up the .tmp on its pre-scan pass
    (void)mgr.listSnapshots();

    assert(!fs::exists(tmp));

    fs::remove_all(dir);
    std::cout << "PASS: orphaned .tmp cleaned by listSnapshots\n";
}

// ──────────────────────────────────────────────────────────
// Test 6: loadWithFallback throws when all snapshots corrupt
// ──────────────────────────────────────────────────────────
void test_loadWithFallback_throws_when_all_corrupt() {
    auto dir = makeSnapshotDir("allcorrupt");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    scfg.maxSnapshots = 10;
    SnapshotManager mgr(scfg);

    MemoryStore store(defaultStoreConfig());
    store.start();
    store.createCollection("docs", CollectionOptions{});
    store.insert("docs", makeDoc("d1"));
    auto s1 = mgr.createSnapshot(store, 1);
    std::this_thread::sleep_for(std::chrono::seconds(1));
    store.insert("docs", makeDoc("d2"));
    auto s2 = mgr.createSnapshot(store, 2);
    store.stop();

    // Corrupt both
    fs::resize_file(s1, fs::file_size(s1) / 2);
    fs::resize_file(s2, fs::file_size(s2) / 2);

    MemoryStore store2(defaultStoreConfig());
    store2.start();
    bool threw = false;
    try {
        mgr.loadWithFallback(store2);
    } catch (const std::exception& e) {
        threw = true;
        std::string msg = e.what();
        assert(msg.find("all") != std::string::npos);
    }
    assert(threw);
    store2.stop();

    fs::remove_all(dir);
    std::cout << "PASS: loadWithFallback throws when all snapshots corrupt\n";
}

// ──────────────────────────────────────────────────────────
// Test 7: v4 chunked format round-trips at scale (multi-chunk path)
// ──────────────────────────────────────────────────────────
void test_v4_format_roundtrip_at_scale() {
    auto dir = makeSnapshotDir("v4scale");

    SnapshotManager::Config scfg;
    scfg.snapshotDir = dir;
    scfg.compressionEnabled = true;
    SnapshotManager mgr(scfg);

    MemoryStore::Config storeCfg = defaultStoreConfig();
    storeCfg.maxMemoryBytes = 512ULL * 1024 * 1024;  // room for 10K x ~2 KB docs
    MemoryStore store(storeCfg);
    store.start();
    store.createCollection("docs", CollectionOptions{});
    // Make each doc ~2 KB so 10,000 docs > 16 MB, forcing multi-chunk body.
    std::string bigString(2000, 'x');
    for (int i = 0; i < 10000; i++) {
        Document d;
        d.id = "d" + std::to_string(i);
        d.set_data(nlohmann::json{{"value", bigString}, {"i", i}});
        store.insert("docs", d);
    }

    auto path = mgr.createSnapshot(store, 1);
    store.stop();

    MemoryStore store2(storeCfg);
    store2.start();
    uint64_t seq = mgr.loadSnapshot(path, store2);
    assert(seq == 1);
    auto d = store2.get("docs", "d9999");
    assert(d.has_value());
    assert((*d).data()["i"].get<int>() == 9999);
    store2.stop();

    fs::remove_all(dir);
    std::cout << "PASS: v4 chunked format round-trips at scale (10K docs, 20MB+)\n";
}

// ──────────────────────────────────────────────────────────
// Test 8: WAL sequence anchoring (regression for v1.7.2)
//
// After a snapshot at sequence S, truncateBefore(S) deletes every WAL
// file (because all entries have sequence ≤ S). The next boot's
// wal_->open() finds no files, so sequence_ stays at 0. Subsequent
// appends would produce 1, 2, 3, … which look "older than the snapshot"
// to the boot AFTER that and get silently dropped by recovery's
// `replay(snapSeq)` filter.
//
// PersistenceManager::replayWal now calls wal_->setSequenceFloor(snapSeq)
// to anchor the WAL counter to the snapshot's walSequence — this test
// exercises that anchoring directly on the WAL.
//
// Bug fingerprint on Zoe (2026-05-07): messages from 10:55–11:41 lost
// across the 11:42 restart; same migration applied twice across the
// two boots because the migrations-applied marker was written with
// sequence 1 and discarded by the next recovery.
// ──────────────────────────────────────────────────────────
void test_wal_sequence_floor_anchors_after_truncate() {
    using smartbotic::database::WriteAheadLog;
    auto dir = makeSnapshotDir("wal-floor");

    WriteAheadLog::Config cfg;
    cfg.walDir = dir;

    // Phase 1 — fresh WAL, append 5 entries, snapshot at seq=5,
    // truncate everything ≤ 5. Mirrors the steady-state cycle.
    {
        WriteAheadLog wal(cfg);
        assert(wal.open());
        for (int i = 0; i < 5; ++i) {
            auto entry = WriteAheadLog::makeInsertEntry(
                "test", makeDoc("d" + std::to_string(i)), "test-node");
            uint64_t seq = wal.append(std::move(entry));
            assert(seq == static_cast<uint64_t>(i + 1));
        }
        assert(wal.currentSequence() == 5);
        wal.truncateBefore(5);
        wal.close();
    }
    // After truncate the on-disk WAL files are deleted. The next boot
    // simulates the post-snapshot startup.
    {
        WriteAheadLog wal(cfg);
        assert(wal.open());
        // Without anchoring, sequence_ would be 0 here — that's the bug.
        // The recovery path is responsible for raising the floor; the
        // WAL itself has no way to know about the snapshot.
        assert(wal.currentSequence() == 0);

        // Now anchor to the snapshot's walSequence (5).
        wal.setSequenceFloor(5);
        assert(wal.currentSequence() == 5);

        // setSequenceFloor must be monotonic — calling with a lower
        // value is a no-op.
        wal.setSequenceFloor(3);
        assert(wal.currentSequence() == 5);

        // Subsequent appends produce 6, 7, … — strictly greater than
        // the snapshot's walSequence, so the next boot's recovery
        // (which calls `replay(snapSeq)` with snapSeq=5) will see them.
        auto entry = WriteAheadLog::makeInsertEntry(
            "test", makeDoc("post"), "test-node");
        uint64_t seq = wal.append(std::move(entry));
        assert(seq == 6);
        wal.close();
    }
    // Phase 2 — re-open one more time, verify sequence_ recovers from
    // the new on-disk file at 6 (open() reads the highest entry).
    {
        WriteAheadLog wal(cfg);
        assert(wal.open());
        assert(wal.currentSequence() >= 6);
        wal.close();
    }

    fs::remove_all(dir);
    std::cout << "PASS: WAL sequence floor anchors after snapshot truncate\n";
}

// ──────────────────────────────────────────────────────────
// Test: snapshot must include evicted documents (v1.7.3 fix)
// ──────────────────────────────────────────────────────────
//
// Pre-v1.7.3: serializeStore() walked only `coll->documents` (in-memory),
// silently omitting docs that had been evicted to stubs. Each snapshot
// rotation truncated the WAL past the snapshot's sequence, deleting the
// only remaining copy of those evicted docs — they were lost forever
// across the next restart.
void test_snapshot_includes_evicted_documents() {
    using smartbotic::database::WriteAheadLog;
    auto snapDir = makeSnapshotDir("evict-snap");
    auto walDir  = makeSnapshotDir("evict-wal");

    WriteAheadLog::Config walCfg;
    walCfg.walDir = walDir;
    WriteAheadLog wal(walCfg);
    if (!wal.open()) {
        std::cerr << "FAIL: wal.open() returned false\n";
        std::exit(2);
    }

    // Tiny memory limit so eviction fires under load.
    MemoryStore::Config storeCfg;
    storeCfg.nodeId = "test";
    storeCfg.maxMemoryBytes = 32ULL * 1024;
    storeCfg.evictionThresholdPercent = 50;
    storeCfg.evictionTargetPercent = 30;
    storeCfg.hotWriteFloorMs = 0;     // disable hot-write protection
    storeCfg.evictionChunkSize = 1000;

    MemoryStore store(storeCfg);
    store.start();
    store.createCollection("docs", CollectionOptions{});

    // Wire load callback: scan WAL for the requested doc, same path
    // production uses (database_service.cpp installs a similar lambda
    // backed by PersistenceManager::loadDocument).
    store.setDocumentLoadCallback(
        [&wal](const std::string& coll, const std::string& id,
               uint64_t /*walSeq*/) -> std::optional<Document> {
            std::optional<Document> found;
            wal.replay(0, [&](const smartbotic::database::WalEntry& e) {
                if (e.collection == coll && e.documentId == id && e.data.has_value()) {
                    Document d = Document::fromJson(*e.data);
                    d.id = id;
                    d.collection = coll;
                    found = d;
                }
            });
            return found;
        });

    // Insert N docs, mirroring to WAL so the load callback can resolve them.
    constexpr int N = 100;
    for (int i = 0; i < N; ++i) {
        Document d = makeDoc("doc-" + std::to_string(i));
        {
            auto tree = d.data();
            tree["padding"] = std::string(300, 'x');  // make eviction meaningful
            d.set_data(tree);
        }
        wal.append(WriteAheadLog::makeInsertEntry("docs", d, "test"));
        store.insert("docs", d);
    }

    // Force eviction.
    store.evictDocuments();
    auto evictedIds = store.getEvictedDocumentIds("docs");
    auto inMemoryDocs = store.getAllDocuments("docs");
    assert(!evictedIds.empty());
    assert(evictedIds.size() + inMemoryDocs.size() == static_cast<size_t>(N));

    // Flush WAL so the load callback (which reads via ifstream) sees every
    // entry. In production the WAL is synced periodically by the persistence
    // manager; tests need to force it.
    wal.sync();

    // Take snapshot.
    SnapshotManager::Config scfg;
    scfg.snapshotDir = snapDir;
    SnapshotManager mgr(scfg);
    auto path = mgr.createSnapshot(store, wal.currentSequence());

    // Load into fresh store with a huge memory budget so eviction doesn't
    // fire during verification — we want to see every doc in-memory.
    MemoryStore::Config newCfg = storeCfg;
    newCfg.maxMemoryBytes = 256ULL * 1024 * 1024;
    MemoryStore newStore(newCfg);
    newStore.start();
    mgr.loadSnapshot(path, newStore);

    auto loadedDocs = newStore.getAllDocuments("docs");
    auto loadedEvicted = newStore.getEvictedDocumentIds("docs");
    const size_t totalLoaded = loadedDocs.size() + loadedEvicted.size();

    // Real check (assert is a no-op in Release builds with NDEBUG).
    if (totalLoaded != static_cast<size_t>(N)) {
        std::cerr << "FAIL: snapshot lost docs — loaded "
                  << loadedDocs.size() << " in-memory + "
                  << loadedEvicted.size() << " evicted = "
                  << totalLoaded << " of " << N << " inserted\n";
        std::exit(1);
    }

    std::unordered_set<std::string> loadedIdSet;
    for (const auto& d : loadedDocs) loadedIdSet.insert(d.id);
    for (const auto& id : loadedEvicted) loadedIdSet.insert(id);
    for (const auto& id : evictedIds) {
        if (loadedIdSet.count(id) != 1) {
            std::cerr << "FAIL: evicted doc " << id
                      << " missing from loaded snapshot\n";
            std::exit(1);
        }
    }

    wal.close();
    fs::remove_all(snapDir);
    fs::remove_all(walDir);

    std::cout << "PASS: snapshot includes evicted documents ("
              << inMemoryDocs.size() << " in-memory + "
              << evictedIds.size() << " evicted, all "
              << N << " round-tripped)\n";
}

// ──────────────────────────────────────────────────────────
// Test: eviction stubs carry the real WAL sequence (v1.7.4 fix)
// ──────────────────────────────────────────────────────────
//
// Pre-v1.7.4 the stub stored doc.version (1, 2, 3…) under the misleading
// name `walSequence`. Production load callbacks therefore couldn't use
// it as a `fromSequence` hint and fell back to scanning the entire WAL
// from sequence 0 on every evicted-doc fault — O(WAL_size) per fault,
// O(WAL_size × evicted_count) per eviction tick.
void test_evict_stub_carries_real_wal_sequence() {
    using smartbotic::database::WriteAheadLog;
    auto walDir = makeSnapshotDir("stub-walseq");

    WriteAheadLog::Config walCfg;
    walCfg.walDir = walDir;
    WriteAheadLog wal(walCfg);
    if (!wal.open()) { std::cerr << "wal.open failed\n"; std::exit(2); }

    MemoryStore::Config storeCfg;
    storeCfg.nodeId = "test";
    storeCfg.maxMemoryBytes = 32ULL * 1024;
    storeCfg.evictionThresholdPercent = 50;
    storeCfg.evictionTargetPercent = 30;
    storeCfg.hotWriteFloorMs = 0;
    storeCfg.evictionChunkSize = 1000;

    MemoryStore store(storeCfg);
    store.start();
    store.createCollection("docs", CollectionOptions{});

    // Insert N docs. Mirror to WAL and record the actual sequence on the
    // store — this is what database_service.cpp's persist callback does.
    constexpr int N = 50;
    std::unordered_map<std::string, uint64_t> expectedSeq;
    for (int i = 0; i < N; ++i) {
        Document d = makeDoc("doc-" + std::to_string(i));
        {
            auto tree = d.data();
            tree["padding"] = std::string(300, 'x');
            d.set_data(tree);
        }
        uint64_t seq = wal.append(WriteAheadLog::makeInsertEntry("docs", d, "test"));
        store.insert("docs", d);
        store.recordWalSequence("docs", d.id, seq);
        expectedSeq[d.id] = seq;
    }

    // Re-write a few docs so their `walSequence` should match the LATEST
    // append, not the first one.
    for (int i = 0; i < 5; ++i) {
        Document d = makeDoc("doc-" + std::to_string(i));
        {
            auto tree = d.data();
            tree["padding"] = std::string(400, 'y');
            d.set_data(tree);
        }
        d.version = 2;
        uint64_t seq = wal.append(WriteAheadLog::makeUpdateEntry("docs", d, "test"));
        store.update("docs", d.id, d);
        store.recordWalSequence("docs", d.id, seq);
        expectedSeq[d.id] = seq;
    }

    store.evictDocuments();
    auto evictedIds = store.getEvictedDocumentIds("docs");
    if (evictedIds.empty()) {
        std::cerr << "FAIL: no docs evicted — eviction config too lax\n";
        std::exit(1);
    }

    // Every evicted stub must have the real (latest-append) seq, not 0
    // and not the per-doc version (1 or 2).
    for (const auto& id : evictedIds) {
        auto stub = store.getEvictedStub("docs", id);
        if (!stub) {
            std::cerr << "FAIL: no stub for " << id << "\n";
            std::exit(1);
        }
        uint64_t want = expectedSeq[id];
        if (stub->walSequence != want) {
            std::cerr << "FAIL: stub.walSequence=" << stub->walSequence
                      << " for " << id << ", expected " << want
                      << " (real WAL sequence)\n";
            std::exit(1);
        }
        // Sanity: must not equal a per-doc version counter (≤ 2 here).
        if (stub->walSequence <= 2) {
            std::cerr << "FAIL: stub.walSequence=" << stub->walSequence
                      << " looks like a per-doc version, not a WAL seq\n";
            std::exit(1);
        }
    }

    wal.close();
    fs::remove_all(walDir);
    std::cout << "PASS: evict stubs carry real WAL sequence ("
              << evictedIds.size() << " stubs verified, all match "
              << "their last wal.append() return value)\n";
}

int main() {
    test_atomic_write_leaves_no_tmp();
    test_snapshot_round_trip();
    test_verification_catches_truncation();
    test_loadWithFallback_uses_older_when_newest_corrupt();
    test_orphaned_tmp_cleaned_by_listSnapshots();
    test_loadWithFallback_throws_when_all_corrupt();
    test_v4_format_roundtrip_at_scale();
    test_wal_sequence_floor_anchors_after_truncate();
    test_snapshot_includes_evicted_documents();
    test_evict_stub_carries_real_wal_sequence();
    std::cout << "\nAll snapshot durability tests PASSED!\n";
    return 0;
}