#826 Fix retry logic

Author: justcoding121

src/Titanium.Web.Proxy/Exceptions/ServerConnectionException.cs renamed to src/Titanium.Web.Proxy/Exceptions/RetryableServerConnectionException.cs

@@ -3,12 +3,12 @@
 namespace Titanium.Web.Proxy.Exceptions
 {
     /// <summary>
-    /// The server connection was closed upon first read with the new connection from pool.
+    /// The server connection was closed upon first write with the new connection from pool.
     /// Should retry the request with a new connection.
     /// </summary>
-    public class ServerConnectionException : ProxyException
+    public class RetryableServerConnectionException : ProxyException
     {
-        internal ServerConnectionException(string message) : base(message)
+        internal RetryableServerConnectionException(string message) : base(message)
         {
         }
 
@@ -17,9 +17,8 @@ internal ServerConnectionException(string message) : base(message)
         /// </summary>
         /// <param name="message"></param>
         /// <param name="e"></param>
-        internal ServerConnectionException(string message, Exception e) : base(message, e)
+        internal RetryableServerConnectionException(string message, Exception e) : base(message, e)
         {
         }
-
     }
 }

src/Titanium.Web.Proxy/Network/RetryPolicy.cs

@@ -35,14 +35,12 @@ internal async Task<RetryResult> ExecuteAsync(Func<TcpServerConnection, Task<boo
 
             while (true)
             {
+                // setup connection
+                currentConnection ??= await generator();
+
                 try
                 {
-                    // setup connection
-                    currentConnection ??= await generator();
-
-                    // try
                     @continue = await action(currentConnection);
-
                 }
                 catch (Exception ex)
                 {

src/Titanium.Web.Proxy/Network/Streams/HttpServerStream.cs

@@ -30,25 +30,19 @@ internal async ValueTask WriteRequestAsync(Request request, CancellationToken ca
 
         internal async ValueTask<ResponseStatusInfo> ReadResponseStatus(CancellationToken cancellationToken = default)
         {
-            try
-            {
-                string httpStatus = await ReadLineAsync(cancellationToken) ??
-                                    throw new ServerConnectionException("Server connection was closed.");
-
-                if (httpStatus == string.Empty)
-                {
-                    // is this really possible?
-                    httpStatus = await ReadLineAsync(cancellationToken) ??
-                                 throw new ServerConnectionException("Server connection was closed. Response status is empty.");
-                }
+            string httpStatus = await ReadLineAsync(cancellationToken) ??
+                                throw new IOException("Invalid http status code.");
 
-                Response.ParseResponseLine(httpStatus, out var version, out int statusCode, out string description);
-                return new ResponseStatusInfo { Version = version, StatusCode = statusCode, Description = description };
-            }
-            catch (Exception e) when (!(e is ServerConnectionException))
+            if (httpStatus == string.Empty)
             {
-                throw new ServerConnectionException("Server connection was closed. Exception while reading the response status.", e);
+                // is this really possible?
+                httpStatus = await ReadLineAsync(cancellationToken) ??
+                             throw new IOException("Response status is empty.");
             }
+
+            Response.ParseResponseLine(httpStatus, out var version, out int statusCode, out string description);
+            return new ResponseStatusInfo { Version = version, StatusCode = statusCode, Description = description };
+
         }
     }
 }

src/Titanium.Web.Proxy/Network/Streams/HttpStream.cs

@@ -932,7 +932,21 @@ public ValueTask WriteLineAsync(string value, CancellationToken cancellationToke
         internal async Task WriteHeadersAsync(HeaderBuilder headerBuilder, CancellationToken cancellationToken = default)
         {
             var buffer = headerBuilder.GetBuffer();
-            await WriteAsync(buffer.Array, buffer.Offset, buffer.Count, true, cancellationToken);
+
+            try
+            {
+                await WriteAsync(buffer.Array, buffer.Offset, buffer.Count, true, cancellationToken);
+            }
+            catch (IOException e)
+            {
+                //throw this as ServerConnectionException so that RetryPolicy can retry with a new server connection.
+                if (this is HttpServerStream)
+                {
+                    throw new RetryableServerConnectionException("Server connection was closed. Exception while sending request line and headers.", e);
+                }
+
+                throw;
+            }
         }
 
         /// <summary>
@@ -1226,7 +1240,6 @@ private async Task copyBytesToStream(IHttpStreamWriter writer, long count, bool
         protected async ValueTask WriteAsync(RequestResponseBase requestResponse, HeaderBuilder headerBuilder, CancellationToken cancellationToken = default)
         {
             var body = requestResponse.CompressBodyAndUpdateContentLength();
-
             headerBuilder.WriteHeaders(requestResponse.Headers);
             await WriteHeadersAsync(headerBuilder, cancellationToken);
 

src/Titanium.Web.Proxy/ProxyServer.cs

@@ -123,7 +123,7 @@ public ProxyServer(string? rootCertificateName, string? rootCertificateIssuerNam
         /// <summary>
         ///     Number of times to retry upon network failures when connection pool is enabled.
         /// </summary>
-        public int NetworkFailureRetryAttempts { get; set; } = 0;
+        public int NetworkFailureRetryAttempts { get; set; } = 1;
 
         /// <summary>
         ///     Is the proxy currently running?

src/Titanium.Web.Proxy/RequestHandler.cs

@@ -288,8 +288,13 @@ private async Task<RetryResult> handleHttpSessionRequest(SessionEventArgs args,
                     noCache,
                     cancellationToken);
 
-            // for connection pool, retry fails until cache is exhausted.   
-            return await retryPolicy<ServerConnectionException>().ExecuteAsync(async connection =>
+            /// Retry with new connection if the initial stream.WriteAsync call to server fails.
+            /// i.e if request line and headers failed to get send.
+            /// Do not retry after reading data from client stream, 
+            /// because subsequent try will not have data to read from client 
+            /// and will hang at clientStream.ReadAsync call.
+            /// So, throw RetryableServerConnectionException only when we are sure we can retry safely.
+            return await retryPolicy<RetryableServerConnectionException>().ExecuteAsync(async connection =>
             {
                 // set the connection and send request headers
                 args.HttpClient.SetConnection(connection);

tests/Titanium.Web.Proxy.IntegrationTests/NestedProxyTests.cs

@@ -195,14 +195,13 @@ public async Task Nested_Proxy_Farm_With_Connection_Cache_Should_Not_Hang()
             var proxies2 = new List<ProxyServer>();
 
             //create a level 2 upstream proxy farm that forwards to server
-            for (int i = 0; i < 2; i++)
+            for (int i = 0; i < 10; i++)
             {
                 var proxy2 = testSuite.GetProxy();
                 proxy2.ProxyBasicAuthenticateFunc += (_, _, _) =>
                 {
                     return Task.FromResult(true);
                 };
-                proxy2.MaxCachedConnections = 100;
                 proxies2.Add(proxy2);
             }
 
@@ -213,7 +212,6 @@ public async Task Nested_Proxy_Farm_With_Connection_Cache_Should_Not_Hang()
             {
                 var proxy1 = testSuite.GetProxy();
                 var proxy2 = proxies2[rnd.Next() % proxies2.Count];
-                proxy1.MaxCachedConnections = 100;
                 var explicitEndpoint = proxy1.ProxyEndPoints.OfType<ExplicitProxyEndPoint>().First();
                 explicitEndpoint.BeforeTunnelConnectRequest += (_, e) =>
                 {