Auto retry message on connection failure

src/StackExchange.Redis/AssemblyInfoHack.cs

@@ -6,5 +6,6 @@
 [assembly: InternalsVisibleTo("StackExchange.Redis.Server, PublicKey=00240000048000009400000006020000002400005253413100040000010001007791a689e9d8950b44a9a8886baad2ea180e7a8a854f158c9b98345ca5009cdd2362c84f368f1c3658c132b3c0f74e44ff16aeb2e5b353b6e0fe02f923a050470caeac2bde47a2238a9c7125ed7dab14f486a5a64558df96640933b9f2b6db188fc4a820f96dce963b662fa8864adbff38e5b4542343f162ecdc6dad16912fff")]
 [assembly: InternalsVisibleTo("StackExchange.Redis.Tests, PublicKey=00240000048000009400000006020000002400005253413100040000010001007791a689e9d8950b44a9a8886baad2ea180e7a8a854f158c9b98345ca5009cdd2362c84f368f1c3658c132b3c0f74e44ff16aeb2e5b353b6e0fe02f923a050470caeac2bde47a2238a9c7125ed7dab14f486a5a64558df96640933b9f2b6db188fc4a820f96dce963b662fa8864adbff38e5b4542343f162ecdc6dad16912fff")]
 [assembly: InternalsVisibleTo("NRediSearch.Test, PublicKey=00240000048000009400000006020000002400005253413100040000010001007791a689e9d8950b44a9a8886baad2ea180e7a8a854f158c9b98345ca5009cdd2362c84f368f1c3658c132b3c0f74e44ff16aeb2e5b353b6e0fe02f923a050470caeac2bde47a2238a9c7125ed7dab14f486a5a64558df96640933b9f2b6db188fc4a820f96dce963b662fa8864adbff38e5b4542343f162ecdc6dad16912fff")]
-
+// For mocking in tests
+[assembly: InternalsVisibleTo("DynamicProxyGenAssembly2, PublicKey=0024000004800000940000000602000000240000525341310004000001000100c547cac37abd99c8db225ef2f6c8a3602f3b3606cc9891605d02baa56104f4cfc0734aa39b93bf7852f7d9266654753cc297e7d2edfe0bac1cdcf9f717241550e0a7b191195b7667bb4f64bcb8e2121380fd1d9d46ad2d92d2d15605093924cceaf74c4861eff62abf69b9291ed0a340e113be11e6a7d3113e92484cf7045cc7")]
 [assembly: CLSCompliant(true)]

src/StackExchange.Redis/CommandRetry/CommandFailureReason.cs

@@ -0,0 +1,25 @@
+namespace StackExchange.Redis
+{
+    /// <summary>
+    /// The reason a command failed to send or complete.
+    /// </summary>
+    public enum CommandFailureReason
+    {
+        /// <summary>
+        /// No open/valid connection was avaialble to send on - we couldn't even write the command.
+        /// </summary>
+        WriteFailure,
+        /// <summary>
+        /// The message was sent, but we lost the connection and this command in-flight.
+        /// </summary>
+        ConnectionFailure,
+        /// <summary>
+        /// Command has timed out, exceeding the sync or async timeout limits
+        /// </summary>
+        Timeout,
+        /// <summary>
+        /// This command failed again, during a retry
+        /// </summary>
+        RetryFailure,
+    }
+}

src/StackExchange.Redis/CommandRetry/CommandRetryPolicy.cs

@@ -0,0 +1,99 @@
+using System;
+
+namespace StackExchange.Redis
+{
+    /// <summary>
+    /// Policy that determines which commands should be retried upon restoration of a lost connection.
+    /// </summary>
+    public abstract class CommandRetryPolicy
+    {
+        /// <summary>
+        /// Creates a policy instance for a specific multiplexer and its commands.
+        /// </summary>
+        /// <param name="muxer">The muleiplexer this policy is for.</param>
+        protected CommandRetryPolicy(ConnectionMultiplexer muxer) { }
+
+        /// <summary>
+        /// Returns the current length of the retry queue.
+        /// </summary>
+        public abstract int CurrentQueueLength { get; }
+
+        /// <summary>
+        /// Returns whether the current queue is processing (e.g. retrying queued commands).
+        /// </summary>
+        public abstract bool CurrentlyProcessing { get; }
+
+        /// <summary>
+        /// Returns the status of the retry mechanism, e.g. what the queue is doing.
+        /// </summary>
+        public abstract string StatusDescription { get; }
+
+        /// <summary>
+        /// Determines if a message is eligible for retrying at all.
+        /// </summary>
+        /// <param name="message">The message to check eligibility for.</param>
+        /// <returns>True if a message is eligible.</returns>
+        internal static bool IsEligible(Message message)
+        {
+            if ((message.Flags & CommandFlags.NoRetry) != 0
+                || ((message.Flags & CommandFlags.RetryIfNotSent) != 0 && message.Status == CommandStatus.Sent)
+                || message.IsAdmin
+                || message.IsInternalCall)
+            {
+                return false;
+            }
+
+            return true;
+        }
+
+        /// <summary>
+        /// Determines if an xception is eligible for retrying at all.
+        /// </summary>
+        /// <param name="exception">The exception to check eligibility for.</param>
+        /// <returns>True if an exception is eligible.</returns>
+        internal static bool IsEligible(Exception exception) => exception is RedisException;
+
+        /// <summary>
+        /// Tries to queue an eligible command.
+        /// Protected because this isn't called directly - eligibility (above) is checked first by the multiplexer.
+        /// </summary>
+        protected internal abstract bool TryQueue(FailedCommand command);
+
+        /// <summary>
+        /// Called when a heartbeat occurs.
+        /// </summary>
+        public abstract void OnHeartbeat();
+
+        /// <summary>
+        /// Called when a multiplexer reconnects.
+        /// </summary>
+        public abstract void OnReconnect();
+
+        /// <summary>
+        /// Default policy - retry only commands which fail before being sent to the server (alias for <see cref="IfNotSent"/>).
+        /// </summary>
+        /// <returns>An instance of a policy that retries only unsent commands.</returns>
+        public static Func<ConnectionMultiplexer, CommandRetryPolicy> Default => IfNotSent;
+
+        /// <summary>
+        /// Retry all commands.
+        /// </summary>
+        /// <returns>An instance of a retry policy that retries all commands.</returns>
+        public static Func<ConnectionMultiplexer, CommandRetryPolicy> Always
+            => mutex => new DefaultCommandRetryPolicy(mutex, commandStatus => true);
+
+        /// <summary>
+        /// Retry only commands which fail before being sent to the server.
+        /// </summary>
+        /// <returns>An instance of a policy that retries only unsent commands.</returns>
+        public static Func<ConnectionMultiplexer, CommandRetryPolicy> IfNotSent
+            => mutex => new DefaultCommandRetryPolicy(mutex, command => command.Status == CommandStatus.WaitingToBeSent);
+
+        /// <summary>
+        /// Never retry a command.
+        /// </summary>
+        /// <returns>An instance of a retry policy that retries no commands.</returns>
+        public static Func<ConnectionMultiplexer, CommandRetryPolicy> Never
+            => mutex => new NeverCommandRetryPolicy(mutex);
+    }
+}

src/StackExchange.Redis/CommandRetry/DefaultCommandRetryPolicy.cs

@@ -0,0 +1,81 @@
+using System;
+
+namespace StackExchange.Redis
+{
+    /// <summary>
+    /// Command retry policy to determine which commands will be retried after a lost connection is retored
+    /// </summary>
+    public class DefaultCommandRetryPolicy : CommandRetryPolicy
+    {
+        private MessageRetryQueue RetryQueue { get; }
+
+        private readonly Func<FailedCommand, bool> _shouldRetry;
+
+        /// <summary>
+        /// Creates a <see cref="DefaultCommandRetryPolicy"/> for the given <see cref="ConnectionMultiplexer"/>.
+        /// </summary>
+        /// <param name="muxer">The <see cref="ConnectionMultiplexer"/> to handle retries for.</param>
+        /// <param name="shouldRetry">Whether a command should be retried.</param>
+        protected internal DefaultCommandRetryPolicy(ConnectionMultiplexer muxer, Func<FailedCommand, bool> shouldRetry) : base(muxer)
+        {
+            _shouldRetry = shouldRetry;
+            var messageRetryHelper = new MessageRetryHelper(muxer);
+            RetryQueue = new MessageRetryQueue(messageRetryHelper);
+        }
+
+        /// <summary>
+        /// Gets the current length of the retry queue.
+        /// </summary>
+        public override int CurrentQueueLength => RetryQueue.CurrentRetryQueueLength;
+
+        /// <summary>
+        /// Returns whether the current queue is processing (e.g. retrying queued commands).
+        /// </summary>
+        public override bool CurrentlyProcessing => RetryQueue.IsRunning;
+
+        /// <summary>
+        /// Returns whether the current queue is processing (e.g. retrying queued commands).
+        /// </summary>
+        public override string StatusDescription => RetryQueue.StatusDescription;
+
+        /// <summary>
+        /// Tries to queue a message for retry if possible.
+        /// </summary>
+        /// <param name="command">The command to tru queueing (contains the message and exception).</param>
+        /// <returns>True if the message was queued.</returns>
+        /// <remarks>Note that this is internal only - external callers cannot override it to bypass the CanRetry checks.</remarks>
+        protected internal override bool TryQueue(FailedCommand command)
+        {
+            // Sanity check if we should be trying this one
+            if (!_shouldRetry.Invoke(command))
+            {
+                return false;
+            }
+
+            if (RetryQueue.TryHandleFailedCommand(command.Message))
+            {
+                // if this message is a new message set the writetime
+                if (command.Message.GetWriteTime() == 0)
+                {
+                    command.Message.SetEnqueued(null);
+                }
+
+                command.Message.ResetStatusToWaitingToBeSent();
+
+                return true;
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Called on heartbeat, evaluating if anything in queue has timed out and need pruning.
+        /// </summary>
+        public override void OnHeartbeat() => RetryQueue.CheckRetryQueueForTimeouts();
+
+        /// <summary>
+        /// Called on a multiplexer reconnect, to start sending anything in the queue.
+        /// </summary>
+        public override void OnReconnect() => RetryQueue.StartRetryQueueProcessor();
+    }
+}

src/StackExchange.Redis/CommandRetry/FailedCommand.cs

@@ -0,0 +1,45 @@
+using System;
+
+namespace StackExchange.Redis
+{
+    /// <summary>
+    /// Command retry policy to act as a no-op for all commands.
+    /// </summary>
+    public sealed class FailedCommand
+    {
+        /// <summary>
+        /// The original/inner message that failed.
+        /// </summary>
+        internal Message Message;
+
+        /// <summary>
+        /// Status of the command.
+        /// </summary>
+        public CommandStatus Status => Message.Status;
+
+        /// <summary>
+        /// The redis command sent.
+        /// </summary>
+        public string CommandAndKey => Message.CommandAndKey;
+
+        /// <summary>
+        /// The reason this command failed, e.g. no connection, timeout, etc.
+        /// </summary>
+        public CommandFailureReason FailureReason { get; }
+
+        /// <summary>
+        /// The exception that happened to create this failed command.
+        /// </summary>
+        public Exception Exception { get; }
+
+        internal static FailedCommand FromWriteFail(Message message, Exception exception) =>
+            new FailedCommand(message, CommandFailureReason.WriteFailure, exception);
+
+        internal FailedCommand(Message message, CommandFailureReason reason, Exception exception)
+        {
+            Message = message;
+            FailureReason = reason;
+            Exception = exception;
+        }
+    }
+}

src/StackExchange.Redis/CommandRetry/IMessageRetryHelper.cs

@@ -0,0 +1,14 @@
+using System;
+using System.Threading.Tasks;
+
+namespace StackExchange.Redis
+{
+    internal interface IMessageRetryHelper
+    {
+        RedisTimeoutException GetTimeoutException(Message message);
+        bool HasTimedOut(Message message);
+        bool IsEndpointAvailable(Message message);
+        void SetExceptionAndComplete(Message message, Exception ex = null);
+        Task<bool> TryResendAsync(Message message);
+    }
+}

src/StackExchange.Redis/CommandRetry/MessageRetryHelper.cs

@@ -0,0 +1,71 @@
+using System;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace StackExchange.Redis
+{
+    internal class MessageRetryHelper : IMessageRetryHelper
+    {
+        private readonly IInternalConnectionMultiplexer multiplexer;
+
+        public MessageRetryHelper(IInternalConnectionMultiplexer multiplexer)
+        {
+            this.multiplexer = multiplexer;
+        }
+
+        public bool HasTimedOut(Message message)
+        {
+            var timeoutMilliseconds = message.ResultBoxIsAsync ? multiplexer.AsyncTimeoutMilliseconds : multiplexer.TimeoutMilliseconds;
+            int millisecondsTaken = unchecked(Environment.TickCount - message.GetWriteTime());
+            return millisecondsTaken >= timeoutMilliseconds;
+        }
+
+        /// <summary>
+        /// Gets the timeout exception for a message.
+        /// </summary>
+        /// <param name="message">The messae to get a message for</param>
+        /// <returns></returns>
+        /// <remarks>
+        /// Not using ExceptionFactory.Timeout as it can cause deadlock while trying to lock writtenawaiting response queue for GetHeadMessages.
+        /// </remarks>
+        public RedisTimeoutException GetTimeoutException(Message message)
+        {
+            var sb = new StringBuilder();
+            sb.Append("Timeout while waiting for connectionrestore ").Append(message.Command).Append(" (").Append(Format.ToString(multiplexer.TimeoutMilliseconds)).Append("ms)");
+            var ex = new RedisTimeoutException(sb.ToString(), message.Status);
+            return ex;
+        }
+
+        public bool IsEndpointAvailable(Message message) => multiplexer.SelectServer(message) != null;
+
+        /// <summary>
+        /// Tries to re-issue a <see cref="Message"/>.
+        /// </summary>
+        /// <param name="message">The message to re-send.</param>
+        /// <returns>Whether the write was successful.</returns>
+        public async Task<bool> TryResendAsync(Message message)
+        {
+            // Use a specific server if one was specified originally, otherwise auto-select
+            // This is important for things like REPLICAOF we really don't want going to another location
+            var server = message.SpecificServer ?? multiplexer.SelectServer(message);
+            if (server != null)
+            {
+                var result = await server.TryWriteAsync(message).ForAwait();
+
+                if (result != WriteResult.Success)
+                {
+                    var ex = multiplexer.GetException(result, message, server);
+                    SetExceptionAndComplete(message, ex);
+                }
+                return true;
+            }
+            return false;
+        }
+
+        public void SetExceptionAndComplete(Message message, Exception ex = null)
+        {
+            var inner = new RedisConnectionException(ConnectionFailureType.UnableToConnect, "Failed while retrying on connection restore: " + ex.Message, ex);
+            message.SetExceptionAndComplete(inner, null, CommandFailureReason.RetryFailure);
+        }
+    }
+}