Fix connection recovery logs

Author: Havret

src/ArtemisNetClient/AutoRecovering/AutoRecoveringConnection.cs

@@ -66,15 +66,11 @@ private AsyncRetryPolicy<IConnection> CreateConnectionRetryPolicy(IRecoveryPolic
                    }, (result, _, context) =>
                    {
                        var retryCount = context.GetRetryCount();
-                       var endpoint = GetCurrentEndpoint(context);
+                       var endpoint =context.GetEndpoint();
                        if (result.Exception != null)
                        {
                            Log.FailedToEstablishConnection(_logger, endpoint, retryCount, result.Exception);
                        }
-                       else
-                       {
-                           Log.ConnectionEstablished(_logger, endpoint, retryCount);
-                       }
                    });
         }
 
@@ -168,11 +164,20 @@ private Task<IConnection> CreateConnection(CancellationToken cancellationToken)
         {
             var ctx = new Context();
             ctx.SetRetryCount(0);
-            return _connectionRetryPolicy.ExecuteAsync((context, ct) =>
+            return _connectionRetryPolicy.ExecuteAsync(async (context, ct) =>
             {
-                var endpoint = GetCurrentEndpoint(context);
+                int retryCount = context.GetRetryCount();
+                var endpoint = GetNextEndpoint(retryCount);
+                context.SetEndpoint(endpoint);
                 var connectionBuilder = new ConnectionBuilder(_loggerFactory, _messageIdPolicyFactory);
-                return connectionBuilder.CreateAsync(endpoint, ct);
+                var connection = await connectionBuilder.CreateAsync(endpoint, ct);
+
+                if (retryCount > 0)
+                {
+                    Log.ConnectionEstablished(_logger, endpoint, retryCount);
+                }
+                
+                return connection;
             }, ctx, cancellationToken);
         }
 
@@ -248,9 +253,8 @@ private async Task DisposeInnerConnection()
             }
         }
 
-        private Endpoint GetCurrentEndpoint(Context context)
+        private Endpoint GetNextEndpoint(int retryCount)
         {
-            int retryCount = context.GetRetryCount();
             return _endpoints[retryCount % _endpoints.Length];
         }
 

src/ArtemisNetClient/Extensions/PollyContextExtensions.cs

@@ -1,12 +1,11 @@
-using System;
-using Polly;
+using Polly;
 
 namespace ActiveMQ.Artemis.Client
 {
     internal static class PollyContextExtensions
     {
         private static string _retryCountKey = "retryCount";
-        private static string _reconnectDelayKey = "reconnectDelay";
+        private static string _endpointKey = "endpoint";
 
         public static void SetRetryCount(this Context context, int retryCount)
         {
@@ -17,22 +16,15 @@ public static int GetRetryCount(this Context context)
         {
             return (int) context[_retryCountKey];
         }
-
-        public static TimeSpan GetReconnectDelay(this Context context, TimeSpan initialReconnectDelay)
+        
+        public static void SetEndpoint(this Context context, Endpoint endpoint)
         {
-            if (context.TryGetValue(_reconnectDelayKey, out var reconnectDelay))
-            {
-                return (TimeSpan) reconnectDelay;
-            }
-            else
-            {
-                return initialReconnectDelay;
-            }
+            context[_endpointKey] = endpoint;
         }
 
-        public static void SetReconnectDelay(this Context context, TimeSpan reconnectDelay)
+        public static Endpoint GetEndpoint(this Context context)
         {
-            context[_reconnectDelayKey] = reconnectDelay;
+            return (Endpoint) context[_endpointKey];
         }
     }
 }