JiShe.CollectBus/src/JiShe.CollectBus.KafkaProducer/Consumer/ConsumerService.cs

using Confluent.Kafka;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using JiShe.CollectBus.Kafka.Attributes;
using Volo.Abp.DependencyInjection;
using JiShe.CollectBus.Kafka.AdminClient;
using static Confluent.Kafka.ConfigPropertyNames;
using System.Collections.Concurrent;
using System.Text.RegularExpressions;
using NUglify.Html;

namespace JiShe.CollectBus.Kafka.Consumer
{
    public class ConsumerService : IConsumerService, IDisposable
    {
        private readonly ILogger<ConsumerService> _logger;
        private readonly IConfiguration _configuration;
        private readonly ConcurrentDictionary<Type, (object Consumer, CancellationTokenSource CTS)>
            _consumerStore = new();

        public ConsumerService(IConfiguration configuration, ILogger<ConsumerService> logger)
        {
            _configuration = configuration;
            _logger = logger;
        }

        #region private 私有方法

        /// <summary>
        /// 创建消费者
        /// </summary>
        /// <typeparam name="TKey"></typeparam>
        /// <typeparam name="TValue"></typeparam>
        /// <returns></returns>
        private IConsumer<TKey, TValue> CreateConsumer<TKey, TValue>(string? groupId = null) where TKey : notnull where TValue : class
        {
            var config = BuildConsumerConfig(groupId);
            return new ConsumerBuilder<TKey, TValue>(config)
                .SetErrorHandler((_, e) => _logger.LogError($"消费者错误: {e.Reason}"))
                .Build();
        }

        private ConsumerConfig BuildConsumerConfig(string? groupId = null)
        {
            var enableAuth = bool.Parse(_configuration["Kafka:EnableAuthorization"]!);

            var config = new ConsumerConfig
            {
                BootstrapServers = _configuration["Kafka:BootstrapServers"],
                GroupId = groupId ?? "default",
                AutoOffsetReset = AutoOffsetReset.Earliest,
                EnableAutoCommit = false // 禁止AutoCommit
            };

            if (enableAuth)
            {
                config.SecurityProtocol = SecurityProtocol.SaslPlaintext;
                config.SaslMechanism = SaslMechanism.Plain;
                config.SaslUsername = _configuration["Kafka:SaslUserName"];
                config.SaslPassword = _configuration["Kafka:SaslPassword"];
            }

            return config;
        }
        #endregion

        /// <summary>
        /// 订阅消息
        /// </summary>
        /// <typeparam name="TKey"></typeparam>
        /// <typeparam name="TValue"></typeparam>
        /// <param name="topic"></param>
        /// <param name="messageHandler"></param>
        /// <returns></returns>
        public async Task SubscribeAsync<TKey, TValue>(string topic, Func<TKey, TValue, Task<bool>> messageHandler, string? groupId = null) where TKey : notnull where TValue : class
        {
            await SubscribeAsync<TKey, TValue>(new[] { topic }, messageHandler, groupId);
        }


        /// <summary>
        /// 订阅消息
        /// </summary>
        /// <typeparam name="TValue"></typeparam>
        /// <param name="topic"></param>
        /// <param name="messageHandler"></param>
        /// <returns></returns>
        public async Task SubscribeAsync<TValue>(string topic, Func<TValue, Task<bool>> messageHandler, string? groupId = null) where TValue : class
        {
            await SubscribeAsync<TValue>(new[] { topic }, messageHandler,groupId);
        }

        /// <summary>
        /// 订阅消息
        /// </summary>
        /// <typeparam name="TKey"></typeparam>
        /// <typeparam name="TValue"></typeparam>
        /// <param name="topics"></param>
        /// <param name="messageHandler"></param>
        /// <returns></returns>
        public async Task SubscribeAsync<TKey, TValue>(string[] topics, Func<TKey, TValue, Task<bool>> messageHandler, string? groupId = null) where TKey : notnull where TValue : class
        {
            var consumerKey = typeof((TKey, TValue));
            var cts = new CancellationTokenSource();

            var consumer = _consumerStore.GetOrAdd(consumerKey, _ =>
            (
                CreateConsumer<TKey, TValue>(groupId),
                cts
            )).Consumer as IConsumer<TKey, TValue>;

            consumer!.Subscribe(topics);

            _ = Task.Run(async () =>
            {
                while (!cts.IsCancellationRequested)
                {
                    try
                    {
                        var result = consumer.Consume(cts.Token);
                        bool sucess=  await messageHandler(result.Message.Key, result.Message.Value);
                        if (sucess)
                        {
                            consumer.Commit(result); // 手动提交
                        }
                    }
                    catch (ConsumeException ex)
                    {
                        _logger.LogError(ex, $"消息消费失败: {ex.Error.Reason}");
                    }
                }
            });
            await Task.CompletedTask;
        }



        /// <summary>
        /// 订阅消息
        /// </summary>
        /// <typeparam name="TKey"></typeparam>
        /// <typeparam name="TValue"></typeparam>
        /// <param name="topics"></param>
        /// <param name="messageHandler"></param>
        /// <returns></returns>
        public async Task SubscribeAsync<TValue>(string[] topics, Func<TValue, Task<bool>> messageHandler, string? groupId) where TValue : class
        {
            var consumerKey = typeof((Null, TValue));
            var cts = new CancellationTokenSource();

            var consumer = _consumerStore.GetOrAdd(consumerKey, _=>
            (
                CreateConsumer<Null, TValue>(groupId),
                cts
            )).Consumer as IConsumer<Null, TValue>;

            consumer!.Subscribe(topics);

            _ = Task.Run(async () =>
            {
                while (!cts.IsCancellationRequested)
                {
                    try
                    {
                        var result = consumer.Consume(cts.Token);
                        bool sucess = await messageHandler(result.Message.Value);
                        if (sucess)
                            consumer.Commit(result); // 手动提交
                    }
                    catch (ConsumeException ex)
                    {
                        _logger.LogError(ex, $"消息消费失败: {ex.Error.Reason}");
                    }
                }
            });
            await Task.CompletedTask;
        }

        /// <summary>
        /// 取消消息订阅
        /// </summary>
        /// <typeparam name="TKey"></typeparam>
        /// <typeparam name="TValue"></typeparam>
        public void Unsubscribe<TKey, TValue>() where TKey : notnull where TValue : class
        {
            var consumerKey = typeof((TKey, TValue));
            if (_consumerStore.TryRemove(consumerKey, out var entry))
            {
                entry.CTS.Cancel();
                (entry.Consumer as IDisposable)?.Dispose();
                entry.CTS.Dispose();
            }
        }

        /// <summary>
        /// 释放资源
        /// </summary>
        public void Dispose()
        {
            foreach (var entry in _consumerStore.Values)
            {
                entry.CTS.Cancel();
                (entry.Consumer as IDisposable)?.Dispose();
                entry.CTS.Dispose();
            }
            _consumerStore.Clear();
        }
    }
}