git.stella-ops.org/src/StellaOps.Feedser.Storage.Mongo.Tests/AdvisoryStorePerformanceTests.cs

using System.Diagnostics;
using System.Linq;
using System.Threading;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using StellaOps.Feedser.Models;
using StellaOps.Feedser.Storage.Mongo;
using StellaOps.Feedser.Storage.Mongo.Advisories;
using StellaOps.Feedser.Storage.Mongo.Aliases;
using StellaOps.Feedser.Storage.Mongo.Migrations;
using Xunit;
using Xunit.Abstractions;

namespace StellaOps.Feedser.Storage.Mongo.Tests;

[Collection("mongo-fixture")]
public sealed class AdvisoryStorePerformanceTests : IClassFixture<MongoIntegrationFixture>
{
    private const int LargeAdvisoryCount = 30;
    private const int AliasesPerAdvisory = 24;
    private const int ReferencesPerAdvisory = 180;
    private const int AffectedPackagesPerAdvisory = 140;
    private const int VersionRangesPerPackage = 4;
    private const int CvssMetricsPerAdvisory = 24;
    private const int ProvenanceEntriesPerAdvisory = 16;
    private static readonly string LargeSummary = new('A', 128 * 1024);
    private static readonly DateTimeOffset BasePublished = new(2024, 1, 1, 0, 0, 0, TimeSpan.Zero);
    private static readonly DateTimeOffset BaseRecorded = new(2024, 1, 1, 0, 0, 0, TimeSpan.Zero);
    private static readonly TimeSpan TotalBudget = TimeSpan.FromSeconds(28);
    private const double UpsertBudgetPerAdvisoryMs = 500;
    private const double FetchBudgetPerAdvisoryMs = 200;
    private const double FindBudgetPerAdvisoryMs = 200;

    private readonly MongoIntegrationFixture _fixture;
    private readonly ITestOutputHelper _output;

    public AdvisoryStorePerformanceTests(MongoIntegrationFixture fixture, ITestOutputHelper output)
    {
        _fixture = fixture;
        _output = output;
    }

    [Fact]
    public async Task UpsertAndQueryLargeAdvisories_CompletesWithinBudget()
    {
        var databaseName = $"feedser-performance-{Guid.NewGuid():N}";
        var database = _fixture.Client.GetDatabase(databaseName);

        try
        {
            var migrationRunner = new MongoMigrationRunner(
                database,
                Array.Empty<IMongoMigration>(),
                NullLogger<MongoMigrationRunner>.Instance,
                TimeProvider.System);

            var bootstrapper = new MongoBootstrapper(
                database,
                Options.Create(new MongoStorageOptions()),
                NullLogger<MongoBootstrapper>.Instance,
                migrationRunner);
            await bootstrapper.InitializeAsync(CancellationToken.None);

            var aliasStore = new AliasStore(database, NullLogger<AliasStore>.Instance);
            var store = new AdvisoryStore(database, aliasStore, NullLogger<AdvisoryStore>.Instance, TimeProvider.System);
            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(45));

            // Warm up collections (indexes, serialization caches) so perf timings exclude one-time setup work.
            var warmup = CreateLargeAdvisory(-1);
            await store.UpsertAsync(warmup, cts.Token);
            _ = await store.FindAsync(warmup.AdvisoryKey, cts.Token);
            _ = await store.GetRecentAsync(1, cts.Token);

            var advisories = Enumerable.Range(0, LargeAdvisoryCount)
                .Select(CreateLargeAdvisory)
                .ToArray();

            var upsertWatch = Stopwatch.StartNew();
            foreach (var advisory in advisories)
            {
                await store.UpsertAsync(advisory, cts.Token);
            }

            upsertWatch.Stop();
            var upsertPerAdvisory = upsertWatch.Elapsed.TotalMilliseconds / LargeAdvisoryCount;

            var fetchWatch = Stopwatch.StartNew();
            var recent = await store.GetRecentAsync(LargeAdvisoryCount, cts.Token);
            fetchWatch.Stop();
            var fetchPerAdvisory = fetchWatch.Elapsed.TotalMilliseconds / LargeAdvisoryCount;

            Assert.Equal(LargeAdvisoryCount, recent.Count);

            var findWatch = Stopwatch.StartNew();
            foreach (var advisory in advisories)
            {
                var fetched = await store.FindAsync(advisory.AdvisoryKey, cts.Token);
                Assert.NotNull(fetched);
            }

            findWatch.Stop();
            var findPerAdvisory = findWatch.Elapsed.TotalMilliseconds / LargeAdvisoryCount;

            var totalElapsed = upsertWatch.Elapsed + fetchWatch.Elapsed + findWatch.Elapsed;

            _output.WriteLine($"Upserted {LargeAdvisoryCount} large advisories in {upsertWatch.Elapsed} ({upsertPerAdvisory:F2} ms/doc).");
            _output.WriteLine($"Fetched recent advisories in {fetchWatch.Elapsed} ({fetchPerAdvisory:F2} ms/doc).");
            _output.WriteLine($"Looked up advisories individually in {findWatch.Elapsed} ({findPerAdvisory:F2} ms/doc).");
            _output.WriteLine($"Total elapsed {totalElapsed}.");

            Assert.True(upsertPerAdvisory <= UpsertBudgetPerAdvisoryMs, $"Upsert exceeded {UpsertBudgetPerAdvisoryMs} ms per advisory: {upsertPerAdvisory:F2} ms.");
            Assert.True(fetchPerAdvisory <= FetchBudgetPerAdvisoryMs, $"GetRecent exceeded {FetchBudgetPerAdvisoryMs} ms per advisory: {fetchPerAdvisory:F2} ms.");
            Assert.True(findPerAdvisory <= FindBudgetPerAdvisoryMs, $"Find exceeded {FindBudgetPerAdvisoryMs} ms per advisory: {findPerAdvisory:F2} ms.");
            Assert.True(totalElapsed <= TotalBudget, $"Mongo advisory operations exceeded total budget {TotalBudget}: {totalElapsed}.");
        }
        finally
        {
            await _fixture.Client.DropDatabaseAsync(databaseName);
        }
    }

    private static Advisory CreateLargeAdvisory(int index)
    {
        var baseKey = $"ADV-LARGE-{index:D4}";
        var published = BasePublished.AddDays(index);
        var modified = published.AddHours(6);

        var aliases = Enumerable.Range(0, AliasesPerAdvisory)
            .Select(i => $"ALIAS-{baseKey}-{i:D4}")
            .ToArray();

        var provenance = Enumerable.Range(0, ProvenanceEntriesPerAdvisory)
            .Select(i => new AdvisoryProvenance(
                source: i % 2 == 0 ? "nvd" : "vendor",
                kind: i % 3 == 0 ? "normalized" : "enriched",
                value: $"prov-{baseKey}-{i:D3}",
                recordedAt: BaseRecorded.AddDays(i)))
            .ToArray();

        var references = Enumerable.Range(0, ReferencesPerAdvisory)
            .Select(i => new AdvisoryReference(
                url: $"https://vuln.example.com/{baseKey}/ref/{i:D4}",
                kind: i % 2 == 0 ? "advisory" : "article",
                sourceTag: $"tag-{i % 7}",
                summary: $"Reference {baseKey} #{i}",
                provenance: provenance[i % provenance.Length]))
            .ToArray();

        var affectedPackages = Enumerable.Range(0, AffectedPackagesPerAdvisory)
            .Select(i => new AffectedPackage(
                type: i % 3 == 0 ? AffectedPackageTypes.Rpm : AffectedPackageTypes.Deb,
                identifier: $"pkg/{baseKey}/{i:D4}",
                platform: i % 4 == 0 ? "linux/x86_64" : "linux/aarch64",
                versionRanges: Enumerable.Range(0, VersionRangesPerPackage)
                    .Select(r => new AffectedVersionRange(
                        rangeKind: r % 2 == 0 ? "semver" : "evr",
                        introducedVersion: $"1.{index}.{i}.{r}",
                        fixedVersion: $"2.{index}.{i}.{r}",
                        lastAffectedVersion: $"1.{index}.{i}.{r}",
                        rangeExpression: $">=1.{index}.{i}.{r} <2.{index}.{i}.{r}",
                        provenance: provenance[(i + r) % provenance.Length]))
                    .ToArray(),
                statuses: Array.Empty<AffectedPackageStatus>(),
                provenance: new[]
                {
                    provenance[i % provenance.Length],
                    provenance[(i + 3) % provenance.Length],
                }))
            .ToArray();

        var cvssMetrics = Enumerable.Range(0, CvssMetricsPerAdvisory)
            .Select(i => new CvssMetric(
                version: i % 2 == 0 ? "3.1" : "2.0",
                vector: $"CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:{(i % 3 == 0 ? "H" : "L")}",
                baseScore: Math.Max(0, 9.8 - i * 0.2),
                baseSeverity: i % 3 == 0 ? "critical" : "high",
                provenance: provenance[i % provenance.Length]))
            .ToArray();

        return new Advisory(
            advisoryKey: baseKey,
            title: $"Large advisory {baseKey}",
            summary: LargeSummary,
            language: "en",
            published: published,
            modified: modified,
            severity: "critical",
            exploitKnown: index % 2 == 0,
            aliases: aliases,
            references: references,
            affectedPackages: affectedPackages,
            cvssMetrics: cvssMetrics,
            provenance: provenance);
    }
}