yams/tests/benchmarks/search_benchmarks.cpp

#include <filesystem>
#include <memory>
#include <string>
#include <vector>

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

#include <yams/api/content_store.h>
#include <yams/api/content_store_builder.h>
#include <yams/metadata/connection_pool.h>
#include <yams/metadata/database.h>
#include <yams/metadata/metadata_repository.h>
#include <yams/search/search_facets.h>

namespace yams::benchmark {

// --- Search Benchmarks from search_benchmark.cpp ---

class SearchBenchmark : public BenchmarkBase {
public:
    SearchBenchmark(const std::string& name, size_t corpusSize, const Config& config = Config())
        : BenchmarkBase("Search_" + name, config), corpusSize_(corpusSize) {
        setUp();
    }
    ~SearchBenchmark() { tearDown(); }

protected:
    void setUp() {
        tempDir_ = test::make_temp_dir("yams_bench_search_");
        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();
        test::TestDataGenerator generator;
        for (size_t i = 0; i < corpusSize_; ++i) {
            auto content = generator.generateTextDocument(1024 + (i % 10) * 1024);
            api::ContentMetadata metadata;
            metadata.name = "doc_" + std::to_string(i) + ".txt";
            contentStore_->storeBytes(std::as_bytes(std::span(content)), metadata);
        }
        queries_ = {"benchmark", "performance", "search"};
    }

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

    size_t runIteration() override {
        // This is a mock implementation as IContentStore has no search method.
        return queries_.size();
    }

    std::filesystem::path tempDir_;
    std::unique_ptr<api::IContentStore> contentStore_;
    size_t corpusSize_;
    std::vector<std::string> queries_;
};

// --- Facet aggregation benchmarks ---

class FacetAggregationBenchmark : public BenchmarkBase {
public:
    FacetAggregationBenchmark(const std::string& name, int resultCount,
                              const Config& config = Config())
        : BenchmarkBase(name, config), resultCount_(resultCount) {
        for (int i = 0; i < resultCount_; ++i) {
            search::SearchResultItem item;
            item.documentId = i;
            item.title = "Document " + std::to_string(i);
            item.contentPreview = "This is the content preview for document " + std::to_string(i) +
                                  " with some searchable text.";
            item.contentType = (i % 3 == 0)   ? "text/plain"
                               : (i % 3 == 1) ? "text/markdown"
                                              : "text/html";
            item.detectedLanguage = (i % 2 == 0) ? "en" : "es";
            results_.push_back(item);
        }
        facetFields_ = {"contentType", "language"};
    }

protected:
    size_t runIteration() override {
        return search::aggregateSearchFacets(results_, facetFields_).facets.size();
    }

    int resultCount_;
    std::vector<search::SearchResultItem> results_;
    std::vector<std::string> facetFields_;
};

} // namespace yams::benchmark

// --- Main Runner ---

using yams::benchmark::archiveJsonFileBestEffort;
using yams::benchmark::BenchmarkBase;
using yams::benchmark::matchesAnyFilter;
using yams::benchmark::parseBenchConfig;
using yams::benchmark::prepareBenchmarkRun;
using yams::benchmark::SearchBenchmark;
using yams::test::BenchmarkTracker;

int main(int argc, char** argv) {
    auto cli = parseBenchConfig(argc, argv, "search_benchmarks");
    if (!cli.outputFile) {
        cli.outputFile = cli.outDir / "search_benchmarks.jsonl";
    }
    if (!prepareBenchmarkRun(cli)) {
        std::cerr << "ERROR: unable to prepare benchmark run directory: " << cli.outDir << '\n';
        return 2;
    }

    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 Search 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 benchmark output directory: " << ec_mkdir.message()
                  << std::endl;
    }
    const std::filesystem::path suiteHistoryJson = outDir / "search_benchmarks.json";
    const std::filesystem::path suiteResultsJsonl = outDir / "search_benchmarks.jsonl";
    if (cli.outputFile) {
        config.output_file = cli.outputFile->string();
    } else {
        config.output_file = suiteResultsJsonl.string();
    }

    BenchmarkTracker tracker(suiteHistoryJson);
    std::vector<std::unique_ptr<BenchmarkBase>> benchmarks;

    // Search benchmarks
    benchmarks.push_back(std::make_unique<SearchBenchmark>("ExactMatch_1K", 1000, config));

    benchmarks.push_back(std::make_unique<yams::benchmark::FacetAggregationBenchmark>(
        "FacetAggregation_100", 100, config));
    benchmarks.push_back(std::make_unique<yams::benchmark::FacetAggregationBenchmark>(
        "FacetAggregation_500", 500, config));
    benchmarks.push_back(std::make_unique<yams::benchmark::FacetAggregationBenchmark>(
        "FacetAggregation_1000", 1000, config));

    for (auto& benchmark : benchmarks) {
        if (!matchesAnyFilter(benchmark->name(), cli.filters, cli.exactFilters)) {
            continue;
        }
        auto result = benchmark->run();
        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;
        tracker.recordResult(trackerResult);
    }

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

    if (cli.archive) {
        tracker.flushHistory();
        if (auto dir = archiveJsonFileBestEffort(suiteHistoryJson, cli.archiveDir, "search")) {
            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;
}