git.stella-ops.org/src/TaskRunner/StellaOps.TaskRunner/StellaOps.TaskRunner.Infrastructure/Execution/MongoPackRunLogStore.cs

using MongoDB.Bson;
using MongoDB.Bson.Serialization.Attributes;
using MongoDB.Driver;
using StellaOps.TaskRunner.Core.Configuration;
using StellaOps.TaskRunner.Core.Execution;

namespace StellaOps.TaskRunner.Infrastructure.Execution;

public sealed class MongoPackRunLogStore : IPackRunLogStore
{
    private readonly IMongoCollection<PackRunLogDocument> collection;

    public MongoPackRunLogStore(IMongoDatabase database, TaskRunnerMongoOptions options)
    {
        ArgumentNullException.ThrowIfNull(database);
        ArgumentNullException.ThrowIfNull(options);

        collection = database.GetCollection<PackRunLogDocument>(options.LogsCollection);
        EnsureIndexes(collection);
    }

    public async Task AppendAsync(string runId, PackRunLogEntry entry, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(runId);
        ArgumentNullException.ThrowIfNull(entry);

        var filter = Builders<PackRunLogDocument>.Filter.Eq(document => document.RunId, runId);

        for (var attempt = 0; attempt < 5; attempt++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var last = await collection
                .Find(filter)
                .SortByDescending(document => document.Sequence)
                .FirstOrDefaultAsync(cancellationToken)
                .ConfigureAwait(false);

            var nextSequence = last is null ? 1 : last.Sequence + 1;

            var document = PackRunLogDocument.FromDomain(runId, nextSequence, entry);

            try
            {
                await collection.InsertOneAsync(document, cancellationToken: cancellationToken).ConfigureAwait(false);
                return;
            }
            catch (MongoWriteException ex) when (ex.WriteError?.Category == ServerErrorCategory.DuplicateKey)
            {
                await Task.Delay(TimeSpan.FromMilliseconds(10), cancellationToken).ConfigureAwait(false);
            }
        }

        throw new InvalidOperationException($"Failed to append log entry for run '{runId}' after multiple attempts.");
    }

    public async IAsyncEnumerable<PackRunLogEntry> ReadAsync(
        string runId,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(runId);

        var filter = Builders<PackRunLogDocument>.Filter.Eq(document => document.RunId, runId);

        using var cursor = await collection
            .Find(filter)
            .SortBy(document => document.Sequence)
            .ToCursorAsync(cancellationToken)
            .ConfigureAwait(false);

        while (await cursor.MoveNextAsync(cancellationToken).ConfigureAwait(false))
        {
            foreach (var document in cursor.Current)
            {
                yield return document.ToDomain();
            }
        }
    }

    public async Task<bool> ExistsAsync(string runId, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(runId);

        var filter = Builders<PackRunLogDocument>.Filter.Eq(document => document.RunId, runId);
        return await collection
            .Find(filter)
            .Limit(1)
            .AnyAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    public static IEnumerable<CreateIndexModel<PackRunLogDocument>> GetIndexModels()
    {
        yield return new CreateIndexModel<PackRunLogDocument>(
            Builders<PackRunLogDocument>.IndexKeys
                .Ascending(document => document.RunId)
                .Ascending(document => document.Sequence),
            new CreateIndexOptions { Unique = true, Name = "pack_run_logs_run_sequence" });

        yield return new CreateIndexModel<PackRunLogDocument>(
            Builders<PackRunLogDocument>.IndexKeys
                .Ascending(document => document.RunId)
                .Ascending(document => document.Timestamp),
            new CreateIndexOptions { Name = "pack_run_logs_run_timestamp" });
    }

    private static void EnsureIndexes(IMongoCollection<PackRunLogDocument> target)
        => target.Indexes.CreateMany(GetIndexModels());

    public sealed class PackRunLogDocument
    {
        [BsonId]
        public ObjectId Id { get; init; }

        public string RunId { get; init; } = default!;

        public long Sequence { get; init; }

        public DateTime Timestamp { get; init; }

        public string Level { get; init; } = default!;

        public string EventType { get; init; } = default!;

        public string Message { get; init; } = default!;

        public string? StepId { get; init; }

        public Dictionary<string, string>? Metadata { get; init; }

        public static PackRunLogDocument FromDomain(string runId, long sequence, PackRunLogEntry entry)
            => new()
            {
                Id = ObjectId.GenerateNewId(),
                RunId = runId,
                Sequence = sequence,
                Timestamp = entry.Timestamp.UtcDateTime,
                Level = entry.Level,
                EventType = entry.EventType,
                Message = entry.Message,
                StepId = entry.StepId,
                Metadata = entry.Metadata is null
                    ? null
                    : new Dictionary<string, string>(entry.Metadata, StringComparer.Ordinal)
            };

        public PackRunLogEntry ToDomain()
        {
            IReadOnlyDictionary<string, string>? metadata = Metadata is null
                ? null
                : new Dictionary<string, string>(Metadata, StringComparer.Ordinal);

            return new PackRunLogEntry(
                new DateTimeOffset(Timestamp, TimeSpan.Zero),
                Level,
                EventType,
                Message,
                StepId,
                metadata);
        }
    }
}