PluralKit/PluralKit.Bot/Services/InteractionDispatchService.cs

using System;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;

using NodaTime;

using Serilog;

namespace PluralKit.Bot
{
    public class InteractionDispatchService: IDisposable
    {
        private static readonly Duration DefaultExpiry = Duration.FromMinutes(15);

        private readonly ConcurrentDictionary<Guid, RegisteredInteraction> _handlers = new();
        private readonly CancellationTokenSource _cts = new();
        private readonly IClock _clock;
        private readonly ILogger _logger;
        private readonly Task _cleanupWorker;
        
        public InteractionDispatchService(IClock clock, ILogger logger)
        {
            _clock = clock;
            _logger = logger.ForContext<InteractionDispatchService>();

            _cleanupWorker = CleanupLoop(_cts.Token);
        }

        public async ValueTask<bool> Dispatch(string customId, InteractionContext context)
        {
            if (!Guid.TryParse(customId, out var customIdGuid))
                return false;
            
            if (!_handlers.TryGetValue(customIdGuid, out var handler))
                return false;

            await handler.Callback.Invoke(context);
            return true;
        }

        public void Unregister(string customId)
        {
            if (!Guid.TryParse(customId, out var customIdGuid))
                return;

            _handlers.TryRemove(customIdGuid, out _);
        }

        public string Register(Func<InteractionContext, Task> callback, Duration? expiry = null)
        {
            var key = Guid.NewGuid();
            var handler = new RegisteredInteraction
            {
                Callback = callback, 
                Expiry = _clock.GetCurrentInstant() + (expiry ?? DefaultExpiry)
            };
            
            _handlers[key] = handler;
            return key.ToString();
        }

        private async Task CleanupLoop(CancellationToken ct)
        {
            while (true)
            {
                DoCleanup();
                await Task.Delay(TimeSpan.FromMinutes(1), ct);
            }
        }

        private void DoCleanup()
        {
            var now = _clock.GetCurrentInstant();
            var removedCount = 0;
            foreach (var (key, value) in _handlers.ToArray())
            {
                if (value.Expiry < now)
                {
                    _handlers.TryRemove(key, out _);
                    removedCount++;
                }
            }

            _logger.Debug("Removed {ExpiredInteractions} expired interactions", removedCount);
        }

        private struct RegisteredInteraction
        {
            public Instant Expiry;
            public Func<InteractionContext, Task> Callback;
        }

        public void Dispose()
        {
            _cts.Cancel(); 
            _cts.Dispose();
        }
    }
}