From 2d0572e21fabe9dc4ae4dbfeda3829c3949aee85 Mon Sep 17 00:00:00 2001
From: Juan Martinez Ramirez <juan.martinezramirez@snowflake.com>
Date: Mon, 29 Jul 2024 15:55:59 -0600
Subject: [PATCH] Added mechanism to handle task cancelation for timeouts and
 errors when trying to execute a HTTP request

---
 .../IntegrationTests/SFConnectionIT.cs        | 49 ++++++++++++++++++-
 .../Client/SnowflakeDbConnection.cs           |  7 +--
 Snowflake.Data/Core/HttpUtil.cs               | 12 ++---
 3 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
index 7c17f9bbd..c9d30576c 100644
--- 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
index fc0ba199d..acaa88ba0 100755
--- 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
index 96d502f40..82a436cc3 100755
--- a/Snowflake.Data/Core/HttpUtil.cs
+++ b/Snowflake.Data/Core/HttpUtil.cs
@@ -458,17 +458,15 @@ 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();
 
@@ -476,7 +474,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
 
                     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);