Add prioritized mode (#1011)

* Add prioritized mode for Consumers

* Add support for consumer pinned priority policy

* Fix build

* Refactor Pin ID handling with reusable extension methods

also added 503 handling to fetch

* Test fixes

* Test fixes

Author: mtmk

src/NATS.Client.JetStream/INatsJSConsumer.cs

@@ -81,6 +81,16 @@ IAsyncEnumerable<NatsJSMsg<T>> FetchAsync<T>(
     /// <exception cref="NatsJSApiException">Server responded with an error.</exception>
     ValueTask RefreshAsync(CancellationToken cancellationToken = default);
 
+    /// <summary>
+    /// Unpins this consumer from the current pinned client.
+    /// </summary>
+    /// <param name="group">The priority group name to unpin.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to cancel the API call.</param>
+    /// <returns>A task representing the asynchronous operation.</returns>
+    /// <exception cref="NatsJSException">There was an issue retrieving the response.</exception>
+    /// <exception cref="NatsJSApiException">Server responded with an error.</exception>
+    ValueTask UnpinAsync(string group, CancellationToken cancellationToken = default);
+
     /// <summary>
     /// Consume a set number of messages from the stream using this consumer.
     /// Returns immediately if no messages are available.

src/NATS.Client.JetStream/INatsJSContext.cs

@@ -161,6 +161,20 @@ IAsyncEnumerable<string> ListConsumerNamesAsync(
     /// <remarks>This feature is only available on NATS server v2.11 and later.</remarks>
     ValueTask<bool> ResumeConsumerAsync(string stream, string consumer, CancellationToken cancellationToken = default);
 
+    /// <summary>
+    /// Unpin a consumer from the currently pinned client.
+    /// </summary>
+    /// <param name="stream">Stream name where consumer is associated to.</param>
+    /// <param name="consumer">Consumer name to be unpinned.</param>
+    /// <param name="group">The priority group name to unpin.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to cancel the API call.</param>
+    /// <returns>A task representing the asynchronous operation.</returns>
+    /// <exception cref="NatsJSException">There was an issue retrieving the response.</exception>
+    /// <exception cref="NatsJSApiException">Server responded with an error.</exception>
+    /// <exception cref="ArgumentException">The <paramref name="stream"/> name is invalid.</exception>
+    /// <exception cref="ArgumentNullException">The <paramref name="stream"/> name is <c>null</c>.</exception>
+    ValueTask UnpinConsumerAsync(string stream, string consumer, string group, CancellationToken cancellationToken = default);
+
     /// <summary>
     /// Calls JetStream Account Info API.
     /// </summary>

src/NATS.Client.JetStream/INatsJSNotification.cs

@@ -33,4 +33,18 @@ public record NatsJSMessageSizeExceedsMaxBytesNotification : INatsJSNotification
     public string Name => "Message Size Exceeds MaxBytes";
 }
 
+/// <summary>
+/// Notification sent when the server reports a pin ID mismatch for pinned client priority policy.
+/// </summary>
+/// <remarks>
+/// This notification indicates that the client's pin ID doesn't match what the server expects.
+/// The client should clear its pin ID and retry the request to get a new pin.
+/// </remarks>
+public record NatsJSPinIdMismatchNotification : INatsJSNotification
+{
+    public static readonly NatsJSPinIdMismatchNotification Default = new();
+
+    public string Name => "Pin ID Mismatch";
+}
+
 public record NatsJSProtocolNotification(string Name, int HeaderCode, string HeaderMessageText) : INatsJSNotification;

src/NATS.Client.JetStream/Internal/NatsJSConsume.cs

@@ -31,6 +31,7 @@ internal class NatsJSConsume<TMsg> : NatsSubBase
     private readonly Timer _timer;
     private readonly Task _pullTask;
     private readonly NatsJSNotificationChannel? _notificationChannel;
+    private readonly NatsJSConsumer? _jsConsumer;
 
     private readonly long _maxMsgs;
     private readonly TimeSpan _expires;
@@ -61,6 +62,7 @@ public NatsJSConsume(
         INatsDeserialize<TMsg> serializer,
         NatsSubOpts? opts,
         NatsJSPriorityGroupOpts? priorityGroup,
+        NatsJSConsumer? jsConsumer,
         CancellationToken cancellationToken)
         : base(context.Connection, context.Connection.SubscriptionManager, subject, queueGroup, opts)
     {
@@ -72,6 +74,7 @@ public NatsJSConsume(
         _consumer = consumer;
         _serializer = serializer;
         _priorityGroup = priorityGroup;
+        _jsConsumer = jsConsumer;
 
         if (notificationHandler is { } handler)
         {
@@ -245,6 +248,8 @@ internal override async ValueTask WriteReconnectCommandsAsync(CommandWriter comm
                 Group = _priorityGroup?.Group,
                 MinPending = _priorityGroup?.MinPending ?? 0,
                 MinAckPending = _priorityGroup?.MinAckPending ?? 0,
+                Priority = _priorityGroup?.Priority ?? 0,
+                Id = _jsConsumer?.GetPinId(),
             };
 
             await commandWriter.PublishAsync(
@@ -354,6 +359,12 @@ protected override async ValueTask ReceiveInternalAsync(
                         _notificationChannel?.Notify(NatsJSLeadershipChangeNotification.Default);
                         ResetPending();
                     }
+                    else if (headers.Code == 423)
+                    {
+                        _logger.LogDebug(NatsJSLogEvents.PinIdMismatch, "Pin ID Mismatch");
+                        NatsJSExtensionsInternal.HandlePinIdMismatch(_jsConsumer, _notificationChannel);
+                        ResetPending();
+                    }
                     else if (headers.Code == 503)
                     {
                         _logger.LogDebug(NatsJSLogEvents.NoResponders, "503 no responders");
@@ -400,6 +411,8 @@ protected override async ValueTask ReceiveInternalAsync(
                     _serializer),
                 _context);
 
+            NatsJSExtensionsInternal.TrySetPinIdFromHeaders(msg.Headers, _jsConsumer);
+
             // Stop feeding the user if we are disposed.
             // We need to exit as soon as possible.
             if (Volatile.Read(ref _disposed) == 0)
@@ -479,6 +492,8 @@ private void Pull(string origin, long batch, long maxBytes) => _pullRequests.Wri
             Group = _priorityGroup?.Group,
             MinPending = _priorityGroup?.MinPending ?? 0,
             MinAckPending = _priorityGroup?.MinAckPending ?? 0,
+            Priority = _priorityGroup?.Priority ?? 0,
+            Id = _jsConsumer?.GetPinId(),
         },
         Origin = origin,
     });

src/NATS.Client.JetStream/Internal/NatsJSExtensionsInternal.cs

@@ -4,8 +4,38 @@ namespace NATS.Client.JetStream.Internal;
 
 internal static class NatsJSExtensionsInternal
 {
+    private const string NatsPinIdHeader = "Nats-Pin-Id";
+
     public static long ToNanos(this TimeSpan timeSpan) => (long)(timeSpan.TotalMilliseconds * 1_000_000);
 
+    /// <summary>
+    /// Handles Pin ID mismatch (423) response by clearing the pin ID and notifying.
+    /// </summary>
+    /// <param name="jsConsumer">The consumer to clear the pin ID on.</param>
+    /// <param name="notificationChannel">The notification channel to notify.</param>
+    public static void HandlePinIdMismatch(NatsJSConsumer? jsConsumer, NatsJSNotificationChannel? notificationChannel)
+    {
+        jsConsumer?.SetPinId(null);
+        notificationChannel?.Notify(NatsJSPinIdMismatchNotification.Default);
+    }
+
+    /// <summary>
+    /// Extracts and sets the Pin ID from message headers if present.
+    /// </summary>
+    /// <param name="headers">The message headers.</param>
+    /// <param name="jsConsumer">The consumer to set the pin ID on.</param>
+    public static void TrySetPinIdFromHeaders(NatsHeaders? headers, NatsJSConsumer? jsConsumer)
+    {
+        if (jsConsumer != null && headers != null && headers.TryGetValue(NatsPinIdHeader, out var pinIdValues))
+        {
+            var pinId = pinIdValues.ToString();
+            if (!string.IsNullOrEmpty(pinId))
+            {
+                jsConsumer.SetPinId(pinId);
+            }
+        }
+    }
+
     public static bool HasTerminalJSError(this NatsHeaders headers)
     {
         // terminal codes

src/NATS.Client.JetStream/Internal/NatsJSFetch.cs

@@ -22,6 +22,7 @@ internal class NatsJSFetch<TMsg> : NatsSubBase
     private readonly Timer _hbTimer;
     private readonly Timer _expiresTimer;
     private readonly NatsJSNotificationChannel? _notificationChannel;
+    private readonly NatsJSConsumer? _jsConsumer;
 
     private readonly long _maxMsgs;
     private readonly long _maxBytes;
@@ -47,6 +48,7 @@ public NatsJSFetch(
         INatsDeserialize<TMsg> serializer,
         NatsSubOpts? opts,
         NatsJSPriorityGroupOpts? priorityGroup,
+        NatsJSConsumer? jsConsumer,
         CancellationToken cancellationToken)
         : base(context.Connection, context.Connection.SubscriptionManager, subject, queueGroup, opts)
     {
@@ -57,6 +59,7 @@ public NatsJSFetch(
         _consumer = consumer;
         _serializer = serializer;
         _priorityGroup = priorityGroup;
+        _jsConsumer = jsConsumer;
 
         if (notificationHandler is { } handler)
         {
@@ -191,6 +194,8 @@ internal override async ValueTask WriteReconnectCommandsAsync(CommandWriter comm
             Group = _priorityGroup?.Group,
             MinPending = _priorityGroup?.MinPending ?? 0,
             MinAckPending = _priorityGroup?.MinAckPending ?? 0,
+            Priority = _priorityGroup?.Priority ?? 0,
+            Id = _jsConsumer?.GetPinId(),
         };
 
         await commandWriter.PublishAsync(
@@ -235,6 +240,17 @@ protected override async ValueTask ReceiveInternalAsync(
                     else if (headers is { Code: 100, Message: NatsHeaders.Messages.IdleHeartbeat })
                     {
                     }
+                    else if (headers.Code == 423)
+                    {
+                        _logger.LogDebug(NatsJSLogEvents.PinIdMismatch, "Pin ID Mismatch");
+                        NatsJSExtensionsInternal.HandlePinIdMismatch(_jsConsumer, _notificationChannel);
+                    }
+                    else if (headers.Code == 503)
+                    {
+                        _logger.LogDebug(NatsJSLogEvents.NoResponders, "503 no responders");
+                        _notificationChannel?.Notify(NatsJSNoRespondersNotification.Default);
+                        EndSubscription(NatsSubEndReason.None);
+                    }
                     else if (headers.HasTerminalJSError())
                     {
                         _userMsgs.Writer.TryComplete(new NatsJSProtocolException(headers.Code, headers.Message, headers.MessageText));
@@ -273,6 +289,8 @@ protected override async ValueTask ReceiveInternalAsync(
                     _serializer),
                 _context);
 
+            NatsJSExtensionsInternal.TrySetPinIdFromHeaders(msg.Headers, _jsConsumer);
+
             _pendingMsgs--;
             _pendingBytes -= msg.Size;
 

src/NATS.Client.JetStream/Models/ConsumerConfigPriorityPolicy.cs

@@ -19,4 +19,9 @@ public enum ConsumerConfigPriorityPolicy
     /// Messages overflow to the next available consumer.
     /// </summary>
     Overflow = 2,
+
+    /// <summary>
+    /// Consumer is pinned to a specific client.
+    /// </summary>
+    PinnedClient = 3,
 }

src/NATS.Client.JetStream/Models/ConsumerInfo.cs

@@ -126,4 +126,12 @@ public record ConsumerInfo
     [System.Text.Json.Serialization.JsonPropertyName("push_bound")]
     [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingDefault)]
     public bool PushBound { get; set; }
+
+    /// <summary>
+    /// Information about the currently defined priority groups.
+    /// </summary>
+    /// <remarks>This feature is only available on NATS server v2.11 and later.</remarks>
+    [System.Text.Json.Serialization.JsonPropertyName("priority_groups")]
+    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingDefault)]
+    public ICollection<PriorityGroupState>? PriorityGroups { get; set; }
 }

src/NATS.Client.JetStream/Models/ConsumerUnpinRequest.cs

@@ -0,0 +1,15 @@
+namespace NATS.Client.JetStream.Models;
+
+/// <summary>
+/// A request to the JetStream $JS.API.CONSUMER.UNPIN API
+/// </summary>
+/// <remarks>This feature is only available on NATS server v2.11 and later.</remarks>
+internal record ConsumerUnpinRequest
+{
+    /// <summary>
+    /// The priority group name to unpin.
+    /// </summary>
+    [System.Text.Json.Serialization.JsonPropertyName("group")]
+    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingDefault)]
+    public string? Group { get; set; }
+}

src/NATS.Client.JetStream/Models/ConsumerUnpinResponse.cs

@@ -0,0 +1,11 @@
+namespace NATS.Client.JetStream.Models;
+
+/// <summary>
+/// A response from the JetStream $JS.API.CONSUMER.UNPIN API
+/// </summary>
+internal record ConsumerUnpinResponse
+{
+    [System.Text.Json.Serialization.JsonPropertyName("success")]
+    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.Never)]
+    public bool Success { get; set; }
+}