track/src/TrackCenter.Application/Rabbit/RabbitClient.cs

using Fake.RabbitMQ;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using RabbitMQ.Client;
using RabbitMQ.Client.Events;

namespace TrackCenter.Application.Rabbit;

public class RabbitClient(
    IRabbitMqChannelPool channelPool,
    IServiceScopeFactory serviceScopeFactory,
    ILogger<RabbitClient> logger,
    IOptions<RabbitOptions> options) : ISingletonDependency, IHostedService
{
    public void Subscribe(ConsumeOptions consumeOptions)
    {
        using var channelAccessor = channelPool.Acquire(consumeOptions.Queue);

        if (!channelAccessor.IsNew) return;
        var channel = channelAccessor.Channel;
        var decl = new DefaultDeclaration(channel);

        consumeOptions.Declaration.Invoke(decl);

        if (consumeOptions.FetchCount > 0)
        {
            channel.BasicQos(0, consumeOptions.FetchCount, false);
        }

        var consumer = new AsyncEventingBasicConsumer(channel);
        consumer.Received += (model, ea) =>
        {
            _ = Task.Run(async () =>
            {
                try
                {
                    await using var scope = serviceScopeFactory.CreateAsyncScope();
                    foreach (var handler in scope.ServiceProvider.GetServices<IRabbitHandler>())
                    {
                        if (handler.Enable(consumeOptions))
                        {
                            await handler.Handle(ea);
                        }
                    }
                    channel.BasicAck(ea.DeliveryTag, false);
                }
                catch (Exception ex)
                {
                    logger.LogError($"rabbit on queue({consumeOptions.Queue}) received error: {ex}");
                    channel.BasicNack(ea.DeliveryTag, false, consumeOptions.FailedRequeue);
                }
            });

            return Task.CompletedTask;
        };

        channel.CallbackException += (sender, args) =>
        {
            logger.LogError(args.Exception, "RabbitMQ channel exception");
        };

        channel.BasicConsume(queue: consumeOptions.Queue,
            autoAck: false,
            consumer: consumer);
    }

    public void Publish(string exchange, string routingKey, byte[] data, Action<IBasicProperties>? options = null)
    {
        using var channelAccessor = channelPool.Acquire(routingKey);
        var props = channelAccessor.Channel.CreateBasicProperties();
        options?.Invoke(props);
        channelAccessor.Channel.BasicPublish(exchange, routingKey, props, data);
    }

    public void PublishByConfirmSelect(
        string exchange,
        string routingKey,
        byte[] data,
        Action<IBasicProperties>? options = null
    )
    {
        using var activity = RabbitInstrumentation.ActivitySource.StartActivity("/" + exchange + "/" + routingKey);
        activity?.SetTag("kind", "publish");
        using var channelAccessor = channelPool.Acquire(routingKey);
        var basicProperties = channelAccessor.Channel.CreateBasicProperties();

        // 配置属性
        if (options != null)
            options(basicProperties);
        var model = channelAccessor.Channel;

        // 启用发布确认
        model.ConfirmSelect();
        var readOnlyMemory = (ReadOnlyMemory<byte>) data;
        model.BasicPublish(exchange, routingKey, basicProperties, readOnlyMemory);
        model.WaitForConfirmsOrDie();
    }

    public void PublishByConfirmSelect(
        string exchange,
        string routingKey,
        IEnumerable<byte[]> datas,
        Action<IBasicProperties>? options = null
    )
    {
        using var activity = RabbitInstrumentation.ActivitySource.StartActivity("/" + exchange + "/" + routingKey);
        activity?.SetTag("kind", "publish");
        using var channelAccessor = channelPool.Acquire(routingKey);
        var basicProperties = channelAccessor.Channel.CreateBasicProperties();

        // 配置属性
        if (options != null)
            options(basicProperties);
        var model = channelAccessor.Channel;

        // 启用发布确认
        model.ConfirmSelect();
        foreach (var data in datas)
        {
            var readOnlyMemory = (ReadOnlyMemory<byte>) data;
            model.BasicPublish(exchange, routingKey, basicProperties, readOnlyMemory);
        }

        model.WaitForConfirmsOrDie();
    }

    public void UnSubscribe(string queueName, string exchange, string routingKey, bool deleteQueue = false)
    {
        logger.LogInformation(
            $"UnSubscribe queue: {queueName} exchange: {exchange} routingKey: {routingKey} deleteQueue: {deleteQueue}");
        using var channelAccessor = channelPool.Acquire(queueName);
        if (deleteQueue)
        {
            var passive = channelAccessor.Channel.QueueDeclarePassive(queueName);
            logger.LogInformation($"Queue {queueName} exists, deleting... {passive.MessageCount} messages in queue");
            channelAccessor.Channel.QueueDelete(queueName);
        }

        channelPool.Release(queueName);
    }

    public Task StartAsync(CancellationToken cancellationToken)
    {
        if (options.Value.Declarations.Count > 0)
        {
            using var channelAccessor = channelPool.Acquire();
            var channel = channelAccessor.Channel;
            var decl = new DefaultDeclaration(channel);
            options.Value.Declarations.ForEach(q => q.Invoke(decl));
        }

        foreach (var opt in options.Value.Consumes)
        {
            Subscribe(opt);
        }

        return Task.CompletedTask;
    }

    public Task StopAsync(CancellationToken cancellationToken)
    {
        channelPool.Dispose();
        return Task.CompletedTask;
    }
}