return/continue check in retry

Author: justcoding121

Titanium.Web.Proxy/ExplicitClientHandler.cs

@@ -203,6 +203,12 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
                         {
                             decryptSsl = false;
                         }
+
+                        if(!decryptSsl)
+                        {
+                            await tcpConnectionFactory.Release(prefetchConnectionTask, true);
+                            prefetchConnectionTask = null; 
+                        }
                     }
 
                     if (cancellationTokenSource.IsCancellationRequested)
@@ -250,8 +256,6 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
                                 (buffer, offset, count) => { connectArgs.OnDataSent(buffer, offset, count); },
                                 (buffer, offset, count) => { connectArgs.OnDataReceived(buffer, offset, count); },
                                 connectArgs.CancellationTokenSource, ExceptionFunc);
-
-
                         }
                         finally
                         {

Titanium.Web.Proxy/Network/RetryPolicy.cs

@@ -11,8 +11,6 @@ internal class RetryPolicy<T> where T : Exception
         private readonly TcpConnectionFactory tcpConnectionFactory;
 
         private TcpServerConnection currentConnection;
-        private Exception exception;
-
         private Policy policy;
 
         internal RetryPolicy(int retries, TcpConnectionFactory tcpConnectionFactory)
@@ -26,15 +24,16 @@ internal RetryPolicy(int retries, TcpConnectionFactory tcpConnectionFactory)
         /// <summary>
         ///     Execute and retry the given action until retry number of times.
         /// </summary>
-        /// <param name="action">The action to retry.</param>
+        /// <param name="action">The action to retry with return value specifying whether caller should continue execution.</param>
         /// <param name="generator">The Tcp connection generator to be invoked to get new connection for retry.</param>
         /// <param name="initialConnection">Initial Tcp connection to use.</param>
         /// <returns>Returns the latest connection used and the latest exception if any.</returns>
-        internal async Task<RetryResult> ExecuteAsync(Func<TcpServerConnection, Task> action,
+        internal async Task<RetryResult> ExecuteAsync(Func<TcpServerConnection, Task<bool>> action,
             Func<Task<TcpServerConnection>> generator, TcpServerConnection initialConnection)
         {
             currentConnection = initialConnection;
-            exception = null;
+            Exception exception = null;
+            bool @continue = true;
 
             try
             {
@@ -46,30 +45,32 @@ await policy.ExecuteAsync(async () =>
                     currentConnection = currentConnection as TcpServerConnection ??
                                       await generator();
                     //try
-                    await action(currentConnection);
+                    @continue = await action(currentConnection);
 
                 });
             }
             catch (Exception e) { exception = e; }
 
-            return new RetryResult(currentConnection, exception);
+            return new RetryResult(currentConnection, exception, @continue);
         }
 
         //get the policy
         private Policy getRetryPolicy()
         {
             return Policy.Handle<T>()
                     .RetryAsync(retries,
-                        onRetryAsync: async (ex, i, context) =>
-                        {
-                            if (currentConnection != null)
-                            {
-                                //close connection on error
-                                await tcpConnectionFactory.Release(currentConnection, true);
-                                currentConnection = null;
-                            }
+                        onRetryAsync: onRetry);
+        }
 
-                        });
+        //before retry clear connection
+        private async Task onRetry(Exception ex, int attempt)
+        {
+            if (currentConnection != null)
+            {
+                //close connection on error
+                await tcpConnectionFactory.Release(currentConnection, true);
+                currentConnection = null;
+            }
         }
     }
 
@@ -78,11 +79,13 @@ internal class RetryResult
         internal bool IsSuccess => Exception == null;
         internal TcpServerConnection LatestConnection { get; }
         internal Exception Exception { get; }
+        internal bool Continue { get; }
 
-        internal RetryResult(TcpServerConnection lastConnection, Exception exception)
+        internal RetryResult(TcpServerConnection lastConnection, Exception exception, bool @continue)
         {
             LatestConnection = lastConnection;
             Exception = exception;
+            Continue = @continue;
         }
     }
 }

Titanium.Web.Proxy/RequestHandler.cs

@@ -220,11 +220,12 @@ await handleWebSocketUpgrade(httpCmd, args, request,
                                         response, clientStream, clientStreamWriter,
                                         serverConnection, cancellationTokenSource, cancellationToken);
                                     closeServerConnection = true;
-                                    return;
+                                    return false;
                                 }
 
                                 // construct the web request that we are going to issue on behalf of the client.
                                 await handleHttpSessionRequestInternal(serverConnection, args);
+                                return true;
 
                             }, generator, connection);
 
@@ -237,6 +238,11 @@ await handleWebSocketUpgrade(httpCmd, args, request,
                                 throw result.Exception;
                             }
 
+                            if(!result.Continue)
+                            {
+                                return;
+                            }
+
                             //user requested
                             if (args.WebSession.CloseServerConnection)
                             {