yams/tests/benchmarks/api_benchmarks.cpp

#include <algorithm>
#include <filesystem>
#include <iterator>
#include <thread>
#include <vector>

#include "../common/fixture_manager.h"
#include "../common/test_data_generator.h"
#include "../common/test_helpers_catch2.h"
#include "benchmark_base.h"

#include <yams/api/content_metadata.h>
#include <yams/api/content_store.h>
#include <yams/api/content_store_builder.h>
#include <yams/chunking/chunker.h>
#include <yams/crypto/hasher.h>
#include <yams/metadata/connection_pool.h>
#include <yams/metadata/database.h>
#include <yams/metadata/document_metadata.h>
#include <yams/metadata/metadata_repository.h>
#include <yams/storage/storage_engine.h>

#include <yams/core/types.h>

using namespace yams;
using namespace yams::benchmark;

// --- Ingestion Benchmarks ---

class IngestionBenchmark : public BenchmarkBase {
public:
    IngestionBenchmark(const std::string& name, const Config& config = Config())
        : BenchmarkBase("Ingestion_" + name, config) {
        setUp();
    }
    ~IngestionBenchmark() { tearDown(); }

protected:
    void setUp() {
        tempDir_ = test::make_temp_dir("yams_bench_ingestion_");
        auto result = api::createContentStore(tempDir_ / "storage");
        if (!result) {
            throw std::runtime_error("Failed to create content store: " + result.error().message);
        }
        contentStore_ = std::move(result).value();
        generator_ = std::make_unique<test::TestDataGenerator>();
    }

    void tearDown() {
        std::error_code ec;
        std::filesystem::remove_all(tempDir_, ec);
    }

    void collectCustomMetrics(std::map<std::string, double>& metrics) override {
        metrics["dedup_ratio"] = lastDedupRatio_;
        metrics["avg_chunk_size"] = avgChunkSize_;
        metrics["total_bytes"] = static_cast<double>(totalBytesProcessed_);
    }

    std::filesystem::path tempDir_;
    std::unique_ptr<api::IContentStore> contentStore_;
    std::unique_ptr<test::TestDataGenerator> generator_;

    double lastDedupRatio_ = 0.0;
    double avgChunkSize_ = 0.0;
    size_t totalBytesProcessed_ = 0;
};

BENCHMARK_F(IngestionBenchmark, SmallDocument) {
    auto content = generator_->generateTextDocument(1024);
    std::vector<std::byte> contentBytes;
    contentBytes.reserve(content.size());
    std::transform(content.begin(), content.end(), std::back_inserter(contentBytes),
                   [](char c) { return std::byte(c); });
    api::ContentMetadata metadata;
    metadata.name = "small_doc.txt";
    auto result = contentStore_->storeBytes(contentBytes, metadata);
    if (result.has_value()) {
        lastDedupRatio_ = result.value().dedupRatio();
        totalBytesProcessed_ = 1024;
        return 1;
    }
    return 0;
}

BENCHMARK_F(IngestionBenchmark, MediumDocument) {
    auto content = generator_->generateTextDocument(100 * 1024);
    std::vector<std::byte> contentBytes;
    contentBytes.reserve(content.size());
    std::transform(content.begin(), content.end(), std::back_inserter(contentBytes),
                   [](char c) { return std::byte(c); });
    api::ContentMetadata metadata;
    metadata.name = "medium_doc.txt";
    auto result = contentStore_->storeBytes(contentBytes, metadata);
    if (result.has_value()) {
        lastDedupRatio_ = result.value().dedupRatio();
        totalBytesProcessed_ = 100 * 1024;
        return 1;
    }
    return 0;
}

// --- Metadata Benchmarks ---

class MetadataBenchmark : public BenchmarkBase {
public:
    MetadataBenchmark(const std::string& name, const Config& config = Config())
        : BenchmarkBase("Metadata_" + name, config) {
        setUp();
    }
    ~MetadataBenchmark() { tearDown(); }

    void collectPragmas(std::map<std::string, std::string>& attrs) {
        if (!connectionPool_) {
            return;
        }

        auto bestEffortPragma = [&](metadata::Database& db, const char* name) {
            auto stmtR = db.prepare(std::string("PRAGMA ") + name);
            if (!stmtR) {
                return;
            }
            auto stmt = std::move(stmtR).value();
            auto stepR = stmt.step();
            if (!stepR || !stepR.value()) {
                return;
            }
            attrs[name] = stmt.getString(0);
        };

        // Best-effort: ignore DB errors; this is for run attribution.
        connectionPool_->withConnection([&](metadata::Database& db) -> yams::Result<void> {
            bestEffortPragma(db, "journal_mode");
            bestEffortPragma(db, "synchronous");
            bestEffortPragma(db, "busy_timeout");
            bestEffortPragma(db, "cache_size");
            bestEffortPragma(db, "temp_store");
            bestEffortPragma(db, "mmap_size");
            bestEffortPragma(db, "foreign_keys");
            bestEffortPragma(db, "wal_autocheckpoint");
            return {};
        });
    }

protected:
    void setUp() {
        tempDir_ = test::make_temp_dir("yams_bench_metadata_");
        auto dbPath = tempDir_ / "metadata.db";
        connectionPool_ = std::make_unique<metadata::ConnectionPool>(dbPath.string());
        metadataRepo_ = std::make_unique<metadata::MetadataRepository>(*connectionPool_);
        createTestDocuments();
        generator_ = std::make_unique<test::TestDataGenerator>();
    }

    void tearDown() {
        metadataRepo_.reset();
        connectionPool_.reset();
        std::error_code ec;
        std::filesystem::remove_all(tempDir_, ec);
    }

    void createTestDocuments() {
        for (size_t i = 0; i < 1000; ++i) {
            metadata::DocumentInfo doc;
            doc.fileName = "file_" + std::to_string(i) + ".txt";
            doc.filePath = (tempDir_ / doc.fileName).string();
            doc.sha256Hash = "hash_" + std::to_string(i);
            doc.fileSize = 1024 * (i % 100 + 1);
            doc.createdTime =
                std::chrono::floor<std::chrono::seconds>(std::chrono::system_clock::now());
            auto result = metadataRepo_->insertDocument(doc);
            if (result) {
                documentIds_.push_back(result.value());
            }
        }
    }

    void collectCustomMetrics(std::map<std::string, double>& metrics) override {
        metrics["num_documents"] = static_cast<double>(documentIds_.size());
        metrics["avg_metadata_size"] = avgMetadataSize_;
        metrics["operations_failed"] = static_cast<double>(failedOperations_);
    }

    std::filesystem::path tempDir_;
    std::unique_ptr<metadata::ConnectionPool> connectionPool_;
    std::unique_ptr<metadata::MetadataRepository> metadataRepo_;
    std::unique_ptr<test::TestDataGenerator> generator_;
    std::vector<int64_t> documentIds_;
    double avgMetadataSize_ = 0;
    size_t failedOperations_ = 0;
};

BENCHMARK_F(MetadataBenchmark, SingleUpdate) {
    static size_t docIndex = 0;

    // Amortize timing overhead (high_resolution_clock + loop) by executing several
    // logically-independent single-row updates per benchmark iteration.
    constexpr size_t kOpsPerIteration = 100;
    size_t succeeded = 0;
    for (size_t i = 0; i < kOpsPerIteration; ++i) {
        const auto& docId = documentIds_[docIndex % documentIds_.size()];
        ++docIndex;
        auto result = metadataRepo_->setMetadata(
            docId, "status", metadata::MetadataValue("updated_" + std::to_string(docIndex)));
        if (!result) {
            ++failedOperations_;
            continue;
        }
        ++succeeded;
    }
    return succeeded;
}

BENCHMARK_F(MetadataBenchmark, BulkUpdate) {
    const size_t batchSize = 100;
    auto metadata = generator_->generateMetadata(5);

    // Build batch entries
    std::vector<std::tuple<int64_t, std::string, metadata::MetadataValue>> entries;
    entries.reserve(batchSize * metadata.size());
    for (size_t i = 0; i < batchSize && i < documentIds_.size(); ++i) {
        for (const auto& [key, value] : metadata) {
            entries.emplace_back(documentIds_[i], key, value);
        }
    }

    // Execute batch update
    auto result = metadataRepo_->setMetadataBatch(entries);
    if (!result) {
        failedOperations_ += entries.size();
        return 0;
    }

    avgMetadataSize_ = static_cast<double>(metadata.size());
    return entries.size();
}

// --- Main Runner ---

int main(int argc, char** argv) {
    const auto cli = parseBenchmarkArgs(argc, argv);

    BenchmarkBase::Config config;
    config.verbose = cli.verbose;
    config.warmup_iterations = cli.warmupIterations;
    config.benchmark_iterations = cli.iterations;
    config.track_memory = cli.trackMemory;

    std::cout << "YAMS API Performance Benchmarks\n";
    std::cout << "====================================\n\n";

    std::filesystem::path outDir = cli.outDir;
    std::error_code ec_mkdir;
    std::filesystem::create_directories(outDir, ec_mkdir);
    if (ec_mkdir) {
        std::cerr << "WARNING: unable to create bench_results directory: " << ec_mkdir.message()
                  << std::endl;
    }
    const std::filesystem::path suiteHistoryJson = outDir / "api_benchmarks.json";
    const std::filesystem::path suiteResultsJsonl = outDir / "api_benchmarks.jsonl";
    if (cli.outputFile) {
        config.output_file = cli.outputFile->string();
    } else {
        config.output_file = suiteResultsJsonl.string();
    }
    test::BenchmarkTracker tracker(suiteHistoryJson);

    std::vector<std::unique_ptr<BenchmarkBase>> benchmarks;
    benchmarks.push_back(std::make_unique<SmallDocumentBenchmark>(config));
    benchmarks.push_back(std::make_unique<MediumDocumentBenchmark>(config));
    benchmarks.push_back(std::make_unique<SingleUpdateBenchmark>(config));
    benchmarks.push_back(std::make_unique<BulkUpdateBenchmark>(config));

    for (auto& benchmark : benchmarks) {
        if (!matchesAnyFilter(benchmark->name(), cli.filters)) {
            continue;
        }
        auto result = benchmark->run();
        test::BenchmarkTracker::BenchmarkResult trackerResult;
        trackerResult.name = result.name;
        trackerResult.value = result.duration_ms;
        trackerResult.unit = "ms";
        trackerResult.timestamp = std::chrono::system_clock::now();
        trackerResult.metrics = result.custom_metrics;
        if (auto* md = dynamic_cast<MetadataBenchmark*>(benchmark.get())) {
            md->collectPragmas(trackerResult.attributes);
        }
        tracker.recordResult(trackerResult);
    }

    tracker.generateReport(outDir / "api_benchmark_report.json");
    tracker.generateMarkdownReport(outDir / "api_benchmark_report.md");

    if (cli.archive) {
        // Ensure the history JSON exists on disk before archiving.
        tracker.flushHistory();
        if (auto dir = archiveJsonFileBestEffort(suiteHistoryJson, cli.archiveDir, "api")) {
            std::error_code ec;
            std::filesystem::copy_file(suiteResultsJsonl, *dir / suiteResultsJsonl.filename(),
                                       std::filesystem::copy_options::overwrite_existing, ec);
            tracker.snapshotTo(*dir / "snapshot.json");
        }
    }

    std::cout << "\n====================================\n";
    std::cout << "Benchmark complete. Reports generated.\n";

    return 0;
}