dotnet-cqrs/Svrnty.CQRS.Events/Storage/InMemoryConsumerRegistry.cs

using System;
using Svrnty.CQRS.Events.Abstractions.Subscriptions;
using Svrnty.CQRS.Events.Subscriptions;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Svrnty.CQRS.Events.Abstractions;

namespace Svrnty.CQRS.Events.Storage;

/// <summary>
/// In-memory implementation of <see cref="IConsumerRegistry"/> for tracking active consumers.
/// Uses concurrent collections for thread-safe consumer management.
/// </summary>
/// <remarks>
/// <para>
/// <strong>Thread Safety:</strong>
/// All operations are thread-safe using <see cref="ConcurrentDictionary{TKey,TValue}"/>.
/// </para>
/// <para>
/// <strong>Stale Consumer Cleanup:</strong>
/// Consumers are automatically marked as stale if they don't send heartbeats.
/// Use <see cref="RemoveStaleConsumersAsync"/> periodically to clean up.
/// </para>
/// </remarks>
public class InMemoryConsumerRegistry : IConsumerRegistry
{
    // (subscriptionId, consumerId) -> ConsumerRegistration
    private readonly ConcurrentDictionary<string, ConsumerRegistration> _consumers = new();

    /// <inheritdoc />
    public Task RegisterConsumerAsync(
        string subscriptionId,
        string consumerId,
        Dictionary<string, string>? metadata = null,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));
        if (string.IsNullOrWhiteSpace(consumerId))
            throw new ArgumentException("Consumer ID cannot be null or whitespace.", nameof(consumerId));

        var key = GetKey(subscriptionId, consumerId);
        var now = DateTimeOffset.UtcNow;

        _consumers.AddOrUpdate(
            key,
            // Add new consumer
            _ => new ConsumerRegistration
            {
                SubscriptionId = subscriptionId,
                ConsumerId = consumerId,
                RegisteredAt = now,
                LastHeartbeat = now,
                Metadata = metadata != null ? new Dictionary<string, string>(metadata) : null
            },
            // Update existing consumer (heartbeat)
            (_, existing) =>
            {
                existing.LastHeartbeat = now;
                if (metadata != null)
                {
                    existing.Metadata = new Dictionary<string, string>(metadata);
                }
                return existing;
            });

        return Task.CompletedTask;
    }

    /// <inheritdoc />
    public Task<bool> UnregisterConsumerAsync(
        string subscriptionId,
        string consumerId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));
        if (string.IsNullOrWhiteSpace(consumerId))
            throw new ArgumentException("Consumer ID cannot be null or whitespace.", nameof(consumerId));

        var key = GetKey(subscriptionId, consumerId);
        var removed = _consumers.TryRemove(key, out _);

        return Task.FromResult(removed);
    }

    /// <inheritdoc />
    public Task<List<string>> GetConsumersAsync(
        string subscriptionId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));

        var consumerIds = _consumers.Values
            .Where(c => c.SubscriptionId == subscriptionId)
            .Select(c => c.ConsumerId)
            .ToList();

        return Task.FromResult(consumerIds);
    }

    /// <inheritdoc />
    public Task<List<ConsumerInfo>> GetConsumerInfoAsync(
        string subscriptionId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));

        var consumers = _consumers.Values
            .Where(c => c.SubscriptionId == subscriptionId)
            .Select(c => new ConsumerInfo
            {
                ConsumerId = c.ConsumerId,
                SubscriptionId = c.SubscriptionId,
                RegisteredAt = c.RegisteredAt,
                LastHeartbeat = c.LastHeartbeat,
                Metadata = c.Metadata
            })
            .ToList();

        return Task.FromResult(consumers);
    }

    /// <inheritdoc />
    public Task<bool> HeartbeatAsync(
        string subscriptionId,
        string consumerId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));
        if (string.IsNullOrWhiteSpace(consumerId))
            throw new ArgumentException("Consumer ID cannot be null or whitespace.", nameof(consumerId));

        var key = GetKey(subscriptionId, consumerId);

        if (_consumers.TryGetValue(key, out var registration))
        {
            registration.LastHeartbeat = DateTimeOffset.UtcNow;
            return Task.FromResult(true);
        }

        return Task.FromResult(false);
    }

    /// <inheritdoc />
    public Task<bool> IsConsumerActiveAsync(
        string subscriptionId,
        string consumerId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(subscriptionId))
            throw new ArgumentException("Subscription ID cannot be null or whitespace.", nameof(subscriptionId));
        if (string.IsNullOrWhiteSpace(consumerId))
            throw new ArgumentException("Consumer ID cannot be null or whitespace.", nameof(consumerId));

        var key = GetKey(subscriptionId, consumerId);
        var exists = _consumers.ContainsKey(key);

        return Task.FromResult(exists);
    }

    /// <inheritdoc />
    public Task<int> RemoveStaleConsumersAsync(
        TimeSpan timeout,
        CancellationToken cancellationToken = default)
    {
        if (timeout <= TimeSpan.Zero)
            throw new ArgumentException("Timeout must be positive.", nameof(timeout));

        var cutoff = DateTimeOffset.UtcNow.Subtract(timeout);
        var staleConsumers = _consumers
            .Where(kvp => kvp.Value.LastHeartbeat < cutoff)
            .Select(kvp => kvp.Key)
            .ToList();

        var removedCount = 0;
        foreach (var key in staleConsumers)
        {
            if (_consumers.TryRemove(key, out _))
            {
                removedCount++;
            }
        }

        return Task.FromResult(removedCount);
    }

    /// <summary>
    /// Get the total number of registered consumers across all subscriptions.
    /// </summary>
    public int GetTotalConsumerCount()
    {
        return _consumers.Count;
    }

    /// <summary>
    /// Get all subscriptions that have at least one active consumer.
    /// </summary>
    public List<string> GetActiveSubscriptions()
    {
        return _consumers.Values
            .Select(c => c.SubscriptionId)
            .Distinct()
            .ToList();
    }

    private static string GetKey(string subscriptionId, string consumerId)
    {
        return $"{subscriptionId}:{consumerId}";
    }

    /// <summary>
    /// Internal consumer registration model.
    /// </summary>
    private class ConsumerRegistration
    {
        public required string SubscriptionId { get; init; }
        public required string ConsumerId { get; init; }
        public required DateTimeOffset RegisteredAt { get; init; }
        public DateTimeOffset LastHeartbeat { get; set; }
        public Dictionary<string, string>? Metadata { get; set; }
    }
}