Add support for initial connection retry

src/NATS.Client.Core/NatsConnection.cs

@@ -37,6 +37,7 @@ public partial class NatsConnection : INatsConnection
     private readonly ILogger<NatsConnection> _logger;
     private readonly ObjectPool _pool;
     private readonly CancellationTokenSource _disposedCts;
+    private readonly CancellationTokenSource _initialConnectCts;
     private readonly string _name;
     private readonly TimeSpan _socketComponentDisposeTimeout = TimeSpan.FromSeconds(5);
     private readonly BoundedChannelOptions _defaultSubscriptionChannelOpts;
@@ -78,6 +79,7 @@ public NatsConnection(NatsOpts opts)
         ConnectionState = NatsConnectionState.Closed;
         _waitForOpenConnection = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
         _disposedCts = new CancellationTokenSource();
+        _initialConnectCts = new CancellationTokenSource();
         _pool = new ObjectPool(opts.ObjectPoolSize);
         _name = opts.Name;
         Counter = new ConnectionStatsCounter();
@@ -170,36 +172,57 @@ internal bool IsDisposed
     /// </summary>
     public async ValueTask ConnectAsync()
     {
-        if (ConnectionState == NatsConnectionState.Open)
-            return;
-
-        TaskCompletionSource? waiter = null;
-        lock (_gate)
+        while (true)
         {
-            ThrowIfDisposed();
-            if (ConnectionState != NatsConnectionState.Closed)
-            {
-                waiter = _waitForOpenConnection;
-            }
-            else
+            try
             {
-                // when closed, change state to connecting and only first connection try-to-connect.
-                ConnectionState = NatsConnectionState.Connecting;
-            }
-        }
+                if (ConnectionState == NatsConnectionState.Open)
+                    return;
 
-        if (waiter != null)
-        {
-            await waiter.Task.ConfigureAwait(false);
-            return;
-        }
+                TaskCompletionSource? waiter = null;
+                lock (_gate)
+                {
+                    ThrowIfDisposed();
+                    if (ConnectionState != NatsConnectionState.Closed)
+                    {
+                        waiter = _waitForOpenConnection;
+                    }
+                    else
+                    {
+                        // when closed, change state to connecting and only first connection try-to-connect.
+                        ConnectionState = NatsConnectionState.Connecting;
+                    }
+                }
 
-        // Only Closed(initial) state, can run initial connect.
-        await InitialConnectAsync().ConfigureAwait(false);
+                if (waiter != null)
+                {
+                    await waiter.Task.ConfigureAwait(false);
+                    return;
+                }
 
-        if (Opts.RequestReplyMode == NatsRequestReplyMode.Direct)
-        {
-            await _subscriptionManager.InitializeInboxSubscriptionAsync(_disposedCts.Token).ConfigureAwait(false);
+                // Only Closed(initial) state, can run initial connect.
+                await InitialConnectAsync().ConfigureAwait(false);
+
+                if (Opts.RequestReplyMode == NatsRequestReplyMode.Direct)
+                {
+                    await _subscriptionManager.InitializeInboxSubscriptionAsync(_disposedCts.Token).ConfigureAwait(false);
+                }
+            }
+            catch (NatsException)
+            {
+                if (!Opts.RetryOnInitialConnect)
+                {
+                    throw;
+                }
+
+                try
+                {
+                    await WaitWithJitterAsync(_initialConnectCts.Token).ConfigureAwait(false);
+                }
+                catch (OperationCanceledException)
+                {
+                }
+            }
         }
     }
 
@@ -254,8 +277,10 @@ public virtual async ValueTask DisposeAsync()
             _waitForOpenConnection.TrySetCanceled();
 #if NET8_0_OR_GREATER
             await _disposedCts.CancelAsync().ConfigureAwait(false);
+            await _initialConnectCts.CancelAsync().ConfigureAwait(false);
 #else
             _disposedCts.Cancel();
+            _initialConnectCts.Cancel();
 #endif
         }
     }
@@ -410,6 +435,9 @@ private async ValueTask InitialConnectAsync()
             _pingTimerCancellationTokenSource = new CancellationTokenSource();
             StartPingTimer(_pingTimerCancellationTokenSource.Token);
             _waitForOpenConnection.TrySetResult();
+#pragma warning disable VSTHRD103
+            _initialConnectCts.Cancel();
+#pragma warning restore VSTHRD103
             _reconnectLoopTask = Task.Run(ReconnectLoop);
             _eventChannel.Writer.TryWrite((NatsEvent.ConnectionOpened, new NatsEventArgs(url?.ToString() ?? string.Empty)));
         }
@@ -713,7 +741,7 @@ private async void ReconnectLoop()
                     _logger.LogDebug(NatsLogEvents.Connection, "Reconnect wait with jitter [{ReconnectCount}]", reconnectCount);
                 }
 
-                await WaitWithJitterAsync().ConfigureAwait(false);
+                await WaitWithJitterAsync(_disposedCts.Token).ConfigureAwait(false);
 
                 if (debug)
                 {
@@ -861,7 +889,7 @@ private NatsUri FixTlsHost(NatsUri uri)
         return uri;
     }
 
-    private async Task WaitWithJitterAsync()
+    private async Task WaitWithJitterAsync(CancellationToken cancellationToken)
     {
         bool stop;
         int retry;
@@ -910,7 +938,7 @@ private async Task WaitWithJitterAsync()
         if (waitTime != TimeSpan.Zero)
         {
             _logger.LogTrace(NatsLogEvents.Connection, "Waiting {WaitMs}ms to reconnect", waitTime.TotalMilliseconds);
-            await Task.Delay(waitTime).ConfigureAwait(false);
+            await Task.Delay(waitTime, cancellationToken).ConfigureAwait(false);
         }
     }
 

src/NATS.Client.Core/NatsOpts.cs

@@ -183,6 +183,29 @@ public sealed record NatsOpts
     /// <seealso cref="INatsTlsUpgradeableSocketConnection"/>
     public INatsSocketConnectionFactory? SocketConnectionFactory { get; init; }
 
+    /// <summary>
+    /// Determines whether the client should retry connecting to the server during the initial connection
+    /// attempt if the connection fails. (default: <c>false</c>)
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This property controls the behavior of the initial connection process.
+    /// </para>
+    /// <para>
+    /// When set to <c>true</c>, the client will periodically retry connecting to the server based on the
+    /// configured retry intervals until a successful connection is established or the operation is explicitly canceled.
+    /// </para>
+    /// <para>
+    /// When set to <c>false</c> (default behavior), the client will not retry and will throw an exception
+    /// if the connection cannot be established initially.
+    /// </para>
+    /// <para>
+    /// The retry intervals are influenced by options such as <see cref="ReconnectWaitMin"/> and <see cref="ReconnectWaitMax"/>, with the potential addition of <see cref="ReconnectJitter"/> for randomized delays between retries.
+    /// This property is particularly useful in scenarios where transient network failures or server unavailability are expected during the first connection attempt.
+    /// </para>
+    /// </remarks>
+    public bool RetryOnInitialConnect { get; init; }
+
     internal NatsUri[] GetSeedUris(bool suppressRandomization = false)
     {
         var urls = Url.Split(',');

tests/NATS.Client.Core.Tests/ConnectionRetryTest.cs

@@ -154,4 +154,29 @@ public async Task Reconnect_doesnt_drop_partially_sent_msgs()
         var loss = 100.0 - (100.0 * received / sent);
         Assert.True(loss <= 1.0, $"message loss of {loss:F}% was above 1% - {sent} sent, {received} received");
     }
+
+    [Fact]
+    public async Task Retry_initial_connect()
+    {
+        await using var server = await NatsServerProcess.StartAsync();
+        await using var connection1 = new NatsConnection(new NatsOpts { Url = server.Url });
+        await using var connection2 = new NatsConnection(new NatsOpts { Url = server.Url, RetryOnInitialConnect = true });
+
+        await server.StopAsync();
+
+        // Preserve the original connection behavior
+        await Assert.ThrowsAsync<NatsException>(async () => await connection1.ConnectAsync());
+
+        var taskStarted = new WaitSignal();
+        var connecting = Task.Run(async () =>
+        {
+            taskStarted.Pulse();
+            await connection2.ConnectAsync();
+        });
+
+        await taskStarted;
+        await server.RestartAsync();
+
+        await connecting;
+    }
 }