git.stella-ops.org/src/Bench/StellaOps.Bench/LinkNotMerge.Vex/StellaOps.Bench.LinkNotMerge.Vex/VexScenarioRunner.cs

using System.Diagnostics;
using EphemeralMongo;
using MongoDB.Bson;
using MongoDB.Driver;

namespace StellaOps.Bench.LinkNotMerge.Vex;

internal sealed class VexScenarioRunner
{
    private readonly VexScenarioConfig _config;
    private readonly IReadOnlyList<VexObservationSeed> _seeds;

    public VexScenarioRunner(VexScenarioConfig config)
    {
        _config = config ?? throw new ArgumentNullException(nameof(config));
        _seeds = VexObservationGenerator.Generate(config);
    }

    public VexScenarioExecutionResult Execute(int iterations, CancellationToken cancellationToken)
    {
        if (iterations <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(iterations), iterations, "Iterations must be positive.");
        }

        var totalDurations = new double[iterations];
        var insertDurations = new double[iterations];
        var correlationDurations = new double[iterations];
        var allocated = new double[iterations];
        var observationThroughputs = new double[iterations];
        var eventThroughputs = new double[iterations];
        VexAggregationResult lastAggregation = new(0, 0, 0, Array.Empty<BsonDocument>());

        for (var iteration = 0; iteration < iterations; iteration++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            using var runner = MongoRunner.Run(new MongoRunnerOptions
            {
                UseSingleNodeReplicaSet = false,
            });

            var client = new MongoClient(runner.ConnectionString);
            var database = client.GetDatabase("linknotmerge_vex_bench");
            var collection = database.GetCollection<BsonDocument>("vex_observations");

            CreateIndexes(collection, cancellationToken);

            var beforeAllocated = GC.GetTotalAllocatedBytes();

            var insertStopwatch = Stopwatch.StartNew();
            InsertObservations(collection, _seeds, _config.ResolveBatchSize(), cancellationToken);
            insertStopwatch.Stop();

            var correlationStopwatch = Stopwatch.StartNew();
            var documents = collection
                .Find(FilterDefinition<BsonDocument>.Empty)
                .Project(Builders<BsonDocument>.Projection
                    .Include("tenant")
                    .Include("statements")
                    .Include("linkset"))
                .ToList(cancellationToken);

            var aggregator = new VexLinksetAggregator();
            lastAggregation = aggregator.Correlate(documents);
            correlationStopwatch.Stop();

            var totalElapsed = insertStopwatch.Elapsed + correlationStopwatch.Elapsed;
            var afterAllocated = GC.GetTotalAllocatedBytes();

            totalDurations[iteration] = totalElapsed.TotalMilliseconds;
            insertDurations[iteration] = insertStopwatch.Elapsed.TotalMilliseconds;
            correlationDurations[iteration] = correlationStopwatch.Elapsed.TotalMilliseconds;
            allocated[iteration] = Math.Max(0, afterAllocated - beforeAllocated) / (1024d * 1024d);

            var totalSeconds = Math.Max(totalElapsed.TotalSeconds, 0.0001d);
            observationThroughputs[iteration] = _seeds.Count / totalSeconds;

            var eventSeconds = Math.Max(correlationStopwatch.Elapsed.TotalSeconds, 0.0001d);
            var eventCount = Math.Max(lastAggregation.EventCount, 1);
            eventThroughputs[iteration] = eventCount / eventSeconds;
        }

        return new VexScenarioExecutionResult(
            totalDurations,
            insertDurations,
            correlationDurations,
            allocated,
            observationThroughputs,
            eventThroughputs,
            ObservationCount: _seeds.Count,
            AliasGroups: _config.ResolveAliasGroups(),
            StatementCount: lastAggregation.StatementCount,
            EventCount: lastAggregation.EventCount,
            AggregationResult: lastAggregation);
    }

    private static void InsertObservations(
        IMongoCollection<BsonDocument> collection,
        IReadOnlyList<VexObservationSeed> seeds,
        int batchSize,
        CancellationToken cancellationToken)
    {
        for (var offset = 0; offset < seeds.Count; offset += batchSize)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var remaining = Math.Min(batchSize, seeds.Count - offset);
            var batch = new List<BsonDocument>(remaining);
            for (var index = 0; index < remaining; index++)
            {
                batch.Add(seeds[offset + index].ToBsonDocument());
            }

            collection.InsertMany(batch, new InsertManyOptions
            {
                IsOrdered = false,
                BypassDocumentValidation = true,
            }, cancellationToken);
        }
    }

    private static void CreateIndexes(IMongoCollection<BsonDocument> collection, CancellationToken cancellationToken)
    {
        var indexKeys = Builders<BsonDocument>.IndexKeys
            .Ascending("tenant")
            .Ascending("linkset.aliases");

        try
        {
            collection.Indexes.CreateOne(new CreateIndexModel<BsonDocument>(indexKeys), cancellationToken: cancellationToken);
        }
        catch
        {
            // non-fatal
        }
    }
}