Add Lame Duck Mode event handler

nats-io · Jan 15, 2025 · ee4710c · ee4710c
1 parent caa4bef
commit ee4710c
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 2 deletions.
diff --git a/src/NATS.Client.Core/INatsConnection.cs b/src/NATS.Client.Core/INatsConnection.cs
@@ -25,6 +25,11 @@ public interface INatsConnection : INatsClient
     /// </summary>
     event AsyncEventHandler<NatsMessageDroppedEventArgs>? MessageDropped;
 
+    /// <summary>
+    /// Event that is raised when server goes into Lame Duck Mode.
+    /// </summary>
+    public event AsyncEventHandler<NatsEventArgs>? LameDuckModeActivated;
+
     /// <summary>
     /// Server information received from the NATS server.
     /// </summary>

diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs
@@ -24,14 +24,13 @@ internal enum NatsEvent
     ConnectionDisconnected,
     ReconnectFailed,
     MessageDropped,
+    LameDuckModeActivated,
 }
 
 public partial class NatsConnection : INatsConnection
 {
 #pragma warning disable SA1401
     internal readonly ConnectionStatsCounter Counter; // allow to call from external sources
-    internal volatile ServerInfo? WritableServerInfo;
-
 #pragma warning restore SA1401
     private readonly object _gate = new object();
     private readonly ILogger<NatsConnection> _logger;
@@ -44,6 +43,7 @@ public partial class NatsConnection : INatsConnection
     private readonly ClientOpts _clientOpts;
     private readonly SubscriptionManager _subscriptionManager;
 
+    private ServerInfo? _writableServerInfo;
     private int _pongCount;
     private int _connectionState;
     private int _isDisposed;
@@ -109,6 +109,8 @@ public NatsConnection(NatsOpts opts)
 
     public event AsyncEventHandler<NatsMessageDroppedEventArgs>? MessageDropped;
 
+    public event AsyncEventHandler<NatsEventArgs>? LameDuckModeActivated;
+
     public INatsConnection Connection => this;
 
     public NatsOpts Opts { get; }
@@ -134,6 +136,21 @@ private set
 
     public Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
 
+    internal ServerInfo? WritableServerInfo
+    {
+        get => Interlocked.CompareExchange(ref _writableServerInfo, null, null);
+        set
+        {
+            var current = Interlocked.CompareExchange(ref _writableServerInfo, null, null);
+            if (current?.LameDuckMode == false && value?.LameDuckMode == true)
+            {
+                _eventChannel.Writer.TryWrite((NatsEvent.LameDuckModeActivated, new NatsEventArgs(string.Empty)));
+            }
+
+            Interlocked.Exchange(ref _writableServerInfo, value);
+        }
+    }
+
     internal bool IsDisposed
     {
         get => Interlocked.CompareExchange(ref _isDisposed, 0, 0) == 1;
@@ -762,6 +779,9 @@ private async Task PublishEventsAsync()
                     case NatsEvent.MessageDropped when MessageDropped != null && args is NatsMessageDroppedEventArgs error:
                         await MessageDropped.InvokeAsync(this, error).ConfigureAwait(false);
                         break;
+                    case NatsEvent.LameDuckModeActivated when LameDuckModeActivated != null:
+                        await LameDuckModeActivated.InvokeAsync(this, args).ConfigureAwait(false);
+                        break;
                     }
                 }
             }

diff --git a/tests/NATS.Client.JetStream.Tests/NatsJsContextFactoryTest.cs b/tests/NATS.Client.JetStream.Tests/NatsJsContextFactoryTest.cs
@@ -92,6 +92,8 @@ public class MockConnection : INatsConnection
         public event AsyncEventHandler<NatsEventArgs>? ReconnectFailed;
 
         public event AsyncEventHandler<NatsMessageDroppedEventArgs>? MessageDropped;
+
+        public event AsyncEventHandler<NatsEventArgs>? LameDuckModeActivated;
 #pragma warning restore CS0067
 
         public INatsServerInfo? ServerInfo { get; } = null;