step3/api/QuestionsApi/Services/AnswerStatisticsStore.cs

using MongoDB.Driver;
using QuestionsApi.Models;
using System.Collections.Concurrent;

namespace QuestionsApi.Services;

public sealed class AnswerStatisticsStore
{
    private const string StateId = "answer-statistics-v1";
    private static readonly TimeSpan CacheDuration = TimeSpan.FromMinutes(5);

    private readonly IMongoCollection<QuestionAttemptCounter> _counters;
    private readonly IMongoCollection<AnswerStatisticsState> _state;
    private readonly QuestionAttemptStore _attempts;
    private readonly ConcurrentDictionary<string, CachedPercentages> _cache = new();
    private volatile bool _backfillComplete;

    public AnswerStatisticsStore(MongoDatabaseProvider mongo, QuestionAttemptStore attempts)
    {
        _attempts = attempts;
        _counters = mongo.Database.GetCollection<QuestionAttemptCounter>("QuestionAttemptCounters");
        _state = mongo.Database.GetCollection<AnswerStatisticsState>("SystemState");

        var counterIndex = Builders<QuestionAttemptCounter>.IndexKeys
            .Ascending(counter => counter.QuestionExternalId)
            .Ascending(counter => counter.Bucket)
            .Ascending(counter => counter.ChoiceId);
        _counters.Indexes.CreateOne(new CreateIndexModel<QuestionAttemptCounter>(
            counterIndex,
            new CreateIndexOptions { Unique = true }));
    }

    public async Task IncrementAsync(QuestionAttempt attempt)
    {
        if (attempt.IsSkip)
            return;

        var filter = Builders<QuestionAttemptCounter>.Filter.And(
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.QuestionExternalId, attempt.QuestionExternalId),
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.Bucket, attempt.CounterBucket),
            Builders<QuestionAttemptCounter>.Filter.Eq(counter => counter.ChoiceId, attempt.ChoiceId));
        var update = Builders<QuestionAttemptCounter>.Update
            .Inc(counter => counter.Count, 1)
            .Set(counter => counter.UpdatedUtc, DateTime.UtcNow)
            .SetOnInsert(counter => counter.QuestionExternalId, attempt.QuestionExternalId)
            .SetOnInsert(counter => counter.Bucket, attempt.CounterBucket)
            .SetOnInsert(counter => counter.ChoiceId, attempt.ChoiceId);

        try
        {
            await _counters.UpdateOneAsync(filter, update, new UpdateOptions { IsUpsert = true });
        }
        catch (MongoWriteException exception) when (exception.WriteError.Category == ServerErrorCategory.DuplicateKey)
        {
            await _counters.UpdateOneAsync(filter, update);
        }
        _cache.TryRemove(attempt.QuestionExternalId, out _);
    }

    public async Task<Dictionary<string, int>> GetPercentagesForQuestionAsync(string questionExternalId)
    {
        var result = await GetPercentagesByQuestionAsync([questionExternalId]);
        return result.GetValueOrDefault(questionExternalId, []);
    }

    public async Task<Dictionary<string, Dictionary<string, int>>> GetPercentagesByQuestionAsync(
        IEnumerable<string> questionExternalIds)
    {
        var ids = questionExternalIds.Where(id => !string.IsNullOrWhiteSpace(id)).Distinct().ToList();
        var output = new Dictionary<string, Dictionary<string, int>>();
        var missing = new List<string>();
        var now = DateTime.UtcNow;

        foreach (var id in ids)
        {
            if (_cache.TryGetValue(id, out var cached) && cached.ExpiresUtc > now)
                output[id] = new Dictionary<string, int>(cached.Percentages);
            else
                missing.Add(id);
        }

        if (missing.Count == 0)
            return output;

        var backfillComplete = await IsBackfillCompleteAsync();
        if (!backfillComplete)
        {
            var rawPercentages = await _attempts.GetPercentagesByQuestionAsync(missing);
            foreach (var id in missing)
            {
                var percentages = rawPercentages.GetValueOrDefault(id, []);
                output[id] = percentages;
                _cache[id] = new CachedPercentages(percentages, now.Add(CacheDuration));
            }
            return output;
        }

        var totals = await _counters.Aggregate()
            .Match(counter => missing.Contains(counter.QuestionExternalId))
            .Group(
                counter => new { counter.QuestionExternalId, counter.ChoiceId },
                group => new QuestionAttemptCounterTotal
                {
                    Key = new QuestionChoiceCounterKey
                    {
                        QuestionExternalId = group.Key.QuestionExternalId,
                        ChoiceId = group.Key.ChoiceId
                    },
                    Count = group.Sum(counter => counter.Count)
                })
            .ToListAsync();

        var grouped = totals.GroupBy(total => total.Key.QuestionExternalId)
            .ToDictionary(group => group.Key, group => group.ToList());

        foreach (var id in missing)
        {
            Dictionary<string, int> percentages;
            if (grouped.TryGetValue(id, out var questionTotals))
            {
                var total = questionTotals.Sum(item => item.Count);
                percentages = total == 0
                    ? []
                    : questionTotals.ToDictionary(
                        item => item.Key.ChoiceId,
                        item => (int)Math.Round(item.Count * 100.0 / total));
            }
            else
            {
                percentages = [];
            }

            output[id] = percentages;
            _cache[id] = new CachedPercentages(percentages, now.Add(CacheDuration));
        }

        return output;
    }

    public async Task<bool> TryAcquireBackfillLeaseAsync(string owner, TimeSpan duration)
    {
        var now = DateTime.UtcNow;
        var filter = Builders<AnswerStatisticsState>.Filter.And(
            Builders<AnswerStatisticsState>.Filter.Eq(state => state.Id, StateId),
            Builders<AnswerStatisticsState>.Filter.Eq(state => state.Completed, false),
            Builders<AnswerStatisticsState>.Filter.Or(
                Builders<AnswerStatisticsState>.Filter.Lt(state => state.LeaseUntilUtc, now),
                Builders<AnswerStatisticsState>.Filter.Eq(state => state.LeaseOwner, owner)));
        var update = Builders<AnswerStatisticsState>.Update
            .SetOnInsert(state => state.Id, StateId)
            .SetOnInsert(state => state.Completed, false)
            .Set(state => state.LeaseOwner, owner)
            .Set(state => state.LeaseUntilUtc, now.Add(duration))
            .Set(state => state.UpdatedUtc, now);

        try
        {
            var leased = await _state.FindOneAndUpdateAsync(
                filter,
                update,
                new FindOneAndUpdateOptions<AnswerStatisticsState>
                {
                    IsUpsert = true,
                    ReturnDocument = ReturnDocument.After
                });
            return leased?.LeaseOwner == owner;
        }
        catch (MongoWriteException exception) when (exception.WriteError.Category == ServerErrorCategory.DuplicateKey)
        {
            return false;
        }
    }

    public async Task MarkBackfillCompleteAsync(string owner)
    {
        var update = Builders<AnswerStatisticsState>.Update
            .Set(state => state.Completed, true)
            .Set(state => state.LeaseOwner, owner)
            .Set(state => state.UpdatedUtc, DateTime.UtcNow);
        await _state.UpdateOneAsync(state => state.Id == StateId && state.LeaseOwner == owner, update);
        _backfillComplete = true;
    }

    public async Task RebuildAllAsync() => await RebuildAsync(null);

    public async Task RebuildQuestionsAsync(IReadOnlyCollection<string> questionExternalIds)
    {
        if (questionExternalIds.Count > 0)
            await RebuildAsync(questionExternalIds);
    }

    private async Task RebuildAsync(IReadOnlyCollection<string>? questionExternalIds)
    {
        var aggregate = _attempts.Collection.Aggregate().Match(attempt => !attempt.IsSkip);
        if (questionExternalIds is not null)
            aggregate = aggregate.Match(attempt => questionExternalIds.Contains(attempt.QuestionExternalId));

        var totals = await aggregate
            .Group(
                attempt => new { attempt.QuestionExternalId, attempt.CounterBucket, attempt.ChoiceId },
                group => new QuestionAttemptCounterAggregate
                {
                    Key = new QuestionChoiceBucketKey
                    {
                        QuestionExternalId = group.Key.QuestionExternalId,
                        Bucket = group.Key.CounterBucket,
                        ChoiceId = group.Key.ChoiceId
                    },
                    Count = group.Sum(_ => 1L)
                })
            .ToListAsync();

        if (totals.Count > 0)
        {
            var now = DateTime.UtcNow;
            var replacements = totals.Select(total =>
            {
                var counter = new QuestionAttemptCounter
                {
                    QuestionExternalId = total.Key.QuestionExternalId,
                    Bucket = total.Key.Bucket,
                    ChoiceId = total.Key.ChoiceId,
                    Count = total.Count,
                    UpdatedUtc = now
                };
                var filter = Builders<QuestionAttemptCounter>.Filter.And(
                    Builders<QuestionAttemptCounter>.Filter.Eq(item => item.QuestionExternalId, counter.QuestionExternalId),
                    Builders<QuestionAttemptCounter>.Filter.Eq(item => item.Bucket, counter.Bucket),
                    Builders<QuestionAttemptCounter>.Filter.Eq(item => item.ChoiceId, counter.ChoiceId));
                return new ReplaceOneModel<QuestionAttemptCounter>(filter, counter) { IsUpsert = true };
            }).ToList();
            await _counters.BulkWriteAsync(replacements, new BulkWriteOptions { IsOrdered = false });
        }

        if (questionExternalIds is null)
            _cache.Clear();
        else
            foreach (var id in questionExternalIds)
                _cache.TryRemove(id, out _);
    }

    private async Task<bool> IsBackfillCompleteAsync()
    {
        if (_backfillComplete)
            return true;
        _backfillComplete = (await _state.Find(state => state.Id == StateId).FirstOrDefaultAsync())?.Completed == true;
        return _backfillComplete;
    }

    private sealed record CachedPercentages(Dictionary<string, int> Percentages, DateTime ExpiresUtc);
}