Added mechanism to handle task cancelation for timeouts and errors wh…

…en trying to execute a HTTP request
snowflakedb · Jul 29, 2024 · 2d0572e · 2d0572e
1 parent ec0c0fa
commit 2d0572e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 14 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
@@ -164,7 +164,7 @@ public void TestConnectionIsNotMarkedAsOpenWhenWasNotCorrectlyOpenedBefore(bool
         }
 
         [Test]
-        public void TestConnectionIsNotMarkedAsOpenWhenWasNotCorrectlyOpenedWithUsingClause()
+        public void TestConnectionIsNotMarkedAsOpenWhenWfasNotCorrectlyOpenedWithUsingClause()
         {
             for (int i = 0; i < 2; ++i)
             {
@@ -2271,6 +2271,53 @@ public void TestUseMultiplePoolsConnectionPoolByDefault()
             // assert
             Assert.AreEqual(ConnectionPoolType.MultipleConnectionPool, poolVersion);
         }
+
+        [Test]
+        public async Task TestShouldThrowExceptionWhenTaskCanceledWhenOpenAsync()
+        {
+            // act
+            var connectionString = "account=myaccount.servername; " +
+                                      " user=admin; password=password; role=ACCOUNTADMIN;" +
+                                      " DB=TEST_DB; warehouse=COMPUTE_WH;" +
+                                      " useproxy=true;proxyhost=no.such.pro.xy;proxyport=8080";
+
+            using (var connection = new SnowflakeDbConnection())
+            {
+                try{
+                    connection.ConnectionString = connectionString;
+                    Task t = connection.OpenAsync();
+                    t.Wait();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<SnowflakeDbException>(e.InnerException);
+                }
+            }
+        }
+
+        [Test]
+        public async Task TestConnectionAsyncTimeoutWithMaxRetryReached()
+        {
+            var mockRestRequester = new MockRetryUntilRestTimeoutRestRequester()
+            {
+                _forceTimeoutForNonLoginRequestsOnly = false
+            };
+
+            using (DbConnection connection = new MockSnowflakeDbConnection(mockRestRequester))
+            {
+                string maxRetryConnStr = ConnectionString + "maxHttpRetries=8;poolingEnabled=true";
+
+                connection.ConnectionString = maxRetryConnStr;
+                try{
+                    Task t = connection.OpenAsync();
+                    t.Wait();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<TaskCanceledException>(e.InnerException);
+                }
+            }
+        }
     }
 }
 

diff --git a/Snowflake.Data/Client/SnowflakeDbConnection.cs b/Snowflake.Data/Client/SnowflakeDbConnection.cs
@@ -318,11 +318,6 @@ public override Task OpenAsync(CancellationToken cancellationToken)
                            SFError.INTERNAL_ERROR,
                            "Unable to connect");
                     }
-                    else if (previousTask.IsCanceled)
-                    {
-                        _connectionState = ConnectionState.Closed;
-                        logger.Debug("Connection canceled");
-                    }
                     else
                     {
                         // Only continue if the session was opened successfully
@@ -331,7 +326,7 @@ public override Task OpenAsync(CancellationToken cancellationToken)
                         OnSessionEstablished();
                     }
                 },
-                cancellationToken);
+                TaskContinuationOptions.NotOnCanceled);
         }
 
         public Mutex GetArrayBindingMutex()

diff --git a/Snowflake.Data/Core/HttpUtil.cs b/Snowflake.Data/Core/HttpUtil.cs
@@ -458,25 +458,23 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                             return response;
                         }
 
-
+                        if (toThrow is TaskCanceledException)
+                        {
+                            throw toThrow;
+                        }
                         throw new SnowflakeDbException(toThrow, SFError.INTERNAL_ERROR,
                             $"Http request failed and max retry {maxRetryCount} reached");
 
                     }
 
-                    if (childCts.IsCancellationRequested)
-                    {
-                        throw new OperationCanceledException($"cancelado!!!");
-                    }
-
                     // Disposing of the response if not null now that we don't need it anymore
                     response?.Dispose();
 
                     requestMessage.RequestUri = updater.Update(errorReason);
 
                     logger.Debug($"Sleep {backOffInSec} seconds and then retry the request, retryCount: {retryCount}");
 
-                    await Task.Delay(TimeSpan.FromSeconds(backOffInSec), cancellationToken).ConfigureAwait(false);
+                    await Task.Delay(TimeSpan.FromSeconds(1), cancellationToken).ConfigureAwait(false);
                     totalRetryTime += backOffInSec;
 
                     var jitter = GetJitter(backOffInSec);