Add support CustomHeadersBuilder option for NATS. (#1519)

dotnetcore · May 15, 2024 · b0ee9ac · b0ee9ac
1 parent f9091f2
commit b0ee9ac
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/src/DotNetCore.CAP.NATS/CAP.NATSOptions.cs b/src/DotNetCore.CAP.NATS/CAP.NATSOptions.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using NATS.Client;
 using NATS.Client.JetStream;
 
@@ -33,6 +34,11 @@ public class NATSOptions
 
         public Action<ConsumerConfiguration.ConsumerConfigurationBuilder>? ConsumerOptions { get; set; }
 
+        /// <summary>
+        /// If you need to get additional native delivery args, you can use this function to write into <see cref="CapHeader" />.
+        /// </summary>
+        public Func<MsgHandlerEventArgs, IServiceProvider, List<KeyValuePair<string, string>>>? CustomHeadersBuilder { get; set; }
+
         public Func<string, string> NormalizeStreamName { get; set; } = origin => origin.Split('.')[0];
     }
 }
diff --git a/src/DotNetCore.CAP.NATS/NATSConsumerClient.cs b/src/DotNetCore.CAP.NATS/NATSConsumerClient.cs
@@ -20,13 +20,15 @@ internal sealed class NATSConsumerClient : IConsumerClient
         private static readonly SemaphoreSlim ConnectionLock = new SemaphoreSlim(initialCount: 1, maxCount: 1);
 
         private readonly string _groupId;
+        private readonly IServiceProvider _serviceProvider;
         private readonly NATSOptions _natsOptions;
 
         private IConnection? _consumerClient;
 
-        public NATSConsumerClient(string groupId, IOptions<NATSOptions> options)
+        public NATSConsumerClient(string groupId, IOptions<NATSOptions> options, IServiceProvider serviceProvider)
         {
             _groupId = groupId;
+            _serviceProvider = serviceProvider;
             _natsOptions = options.Value ?? throw new ArgumentNullException(nameof(options));
         }
 
@@ -142,6 +144,15 @@ private void SubscriptionMessageHandler(object? sender, MsgHandlerEventArgs e)
 
             headers.Add(Headers.Group, _groupId);
 
+            if (_natsOptions.CustomHeadersBuilder != null)
+            {
+                var customHeaders = _natsOptions.CustomHeadersBuilder(e, _serviceProvider);
+                foreach (var customHeader in customHeaders)
+                {
+                    headers[customHeader.Key] = customHeader.Value;
+                }
+            }
+
             OnMessageCallback!(new TransportMessage(headers, e.Message.Data), e.Message);
         }
 

diff --git a/src/DotNetCore.CAP.NATS/NATSConsumerClientFactory.cs b/src/DotNetCore.CAP.NATS/NATSConsumerClientFactory.cs
@@ -1,6 +1,7 @@
 // Copyright (c) .NET Core Community. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System;
 using DotNetCore.CAP.Transport;
 using Microsoft.Extensions.Options;
 
@@ -9,17 +10,19 @@ namespace DotNetCore.CAP.NATS
     internal sealed class NATSConsumerClientFactory : IConsumerClientFactory
     {
         private readonly IOptions<NATSOptions> _natsOptions;
+        private readonly IServiceProvider _serviceProvider;
 
-        public NATSConsumerClientFactory(IOptions<NATSOptions> natsOptions)
+        public NATSConsumerClientFactory(IOptions<NATSOptions> natsOptions, IServiceProvider serviceProvider)
         {
             _natsOptions = natsOptions;
+            _serviceProvider = serviceProvider;
         }
 
         public IConsumerClient Create(string groupId)
         {
             try
             {
-                var client = new NATSConsumerClient(groupId, _natsOptions);
+                var client = new NATSConsumerClient(groupId, _natsOptions, _serviceProvider);
                 client.Connect();
                 return client;
             }