SNOW-2044393 Fix for returning cancelled sessions to the connection pool (#1150)

Author: sfc-gh-dstempniak

Snowflake.Data.Tests/IntegrationTests/ConnectionMultiplePoolsIT.cs

@@ -440,5 +440,55 @@ private async Task<SnowflakeDbConnection> OpenConnectionAsync(string connectionS
             await connection.OpenAsync().ConfigureAwait(false);
             return connection;
         }
+
+        [Test]
+        public void TestReturningCancelledSessionsToThePool([Values]bool cancelAsync)
+        {
+            var connectionString = ConnectionString + "minPoolSize=0;maxPoolSize=2;application=TestReturningCancelledSessionsToThePool";
+
+            var pool = SnowflakeDbConnectionPool.ConnectionManager.GetPool(connectionString);
+            pool.ClearSessions();
+
+            // pool is empty
+            Assert.AreEqual(0, pool.GetCurrentState().IdleSessionsCount);
+            Assert.AreEqual(0, pool.GetCurrentState().BusySessionsCount);
+
+            var cts = new CancellationTokenSource();
+            var task = Task.Run(async () =>
+            {
+                using (var connection = new SnowflakeDbConnection(connectionString))
+                {
+                    await connection.OpenAsync(cts.Token);
+                    using (var command = connection.CreateCommand())
+                    {
+                        command.CommandText = $"SELECT SYSTEM$WAIT(20)";
+                        await command.ExecuteNonQueryAsync(cts.Token);
+                    }
+                }
+            }, CancellationToken.None);
+
+            Thread.Sleep(2000);
+
+            // one busy session
+            Assert.AreEqual(0, pool.GetCurrentState().IdleSessionsCount);
+            Assert.AreEqual(1, pool.GetCurrentState().BusySessionsCount);
+
+            if (cancelAsync)
+#if NET8_0_OR_GREATER
+                cts.CancelAsync();
+#else
+                cts.Cancel();
+#endif
+            else
+                cts.Cancel();
+
+            // operation cancelled properly
+            var thrown = Assert.Throws<AggregateException>(() => task.Wait());
+            Assert.IsInstanceOf<OperationCanceledException>(thrown.InnerException);
+
+            // one idle session
+            Assert.AreEqual(1, pool.GetCurrentState().IdleSessionsCount);
+            Assert.AreEqual(0, pool.GetCurrentState().BusySessionsCount);
+        }
     }
 }

Snowflake.Data.Tests/IntegrationTests/ConnectionPoolCommonIT.cs

@@ -178,7 +178,7 @@ public void TestTransactionStatusNotTrackedForNonExplicitTransactionCalls()
             }
         }
 
-                [Test]
+        [Test]
         public void TestRollbackTransactionOnPooledWhenConnectionClose()
         {
             var connectionString = SetPoolWithOneElement();
@@ -215,7 +215,6 @@ public void TestRollbackTransactionOnPooledWhenConnectionClose()
             Assert.AreEqual(1, SnowflakeDbConnectionPool.GetCurrentPoolSize(), "Connection should be returned to the pool");
         }
 
-
         private string SetPoolWithOneElement()
         {
             if (_connectionPoolTypeUnderTest == ConnectionPoolType.SingleConnectionCache)
@@ -230,6 +229,5 @@ private int ExpectedPoolCountAfterOpen()
         {
             return _connectionPoolTypeUnderTest == ConnectionPoolType.SingleConnectionCache ? 0 : 1;
         }
-
     }
 }

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

@@ -1950,20 +1950,12 @@ public void TestCancelLoginBeforeTimeout()
                 conn.ConnectionString = infiniteLoginTimeOut;
 
                 Assert.AreEqual(conn.State, ConnectionState.Closed);
-                // At this point the connection string has not been parsed, it will return the
-                // default value
-                //Assert.AreEqual(SFSessionHttpClientProperties.s_retryTimeoutDefault, conn.ConnectionTimeout);
 
                 CancellationTokenSource connectionCancelToken = new CancellationTokenSource();
                 Task connectTask = conn.OpenAsync(connectionCancelToken.Token);
 
-                // Sleep for more than the default timeout to make sure there are no false positive)
-                Thread.Sleep(SFSessionHttpClientProperties.DefaultRetryTimeout.Add(TimeSpan.FromSeconds(10)));
-
                 Assert.AreEqual(ConnectionState.Connecting, conn.State);
 
-                // Cancel the connection because it will never succeed since there is no
-                // connection_timeout defined
                 logger.Debug("connectionCancelToken.Cancel ");
                 connectionCancelToken.Cancel();
 

Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs

@@ -45,7 +45,7 @@ public override void Open()
 
         public override Task OpenAsync(CancellationToken cancellationToken)
         {
-            registerConnectionCancellationCallback(cancellationToken);
+            cancellationToken.Register(() => { _connectionState = ConnectionState.Closed; });
 
             SetMockSession();
 

Snowflake.Data/Client/SnowflakeDbConnection.cs

@@ -314,7 +314,6 @@ public override Task OpenAsync(CancellationToken cancellationToken)
                 logger.Debug($"Open with a connection already opened: {_connectionState}");
                 return Task.CompletedTask;
             }
-            registerConnectionCancellationCallback(cancellationToken);
             OnSessionConnecting();
             FillConnectionStringFromTomlConfigIfNotSet();
             return SnowflakeDbConnectionPool
@@ -423,20 +422,6 @@ protected override void Dispose(bool disposing)
             base.Dispose(disposing);
         }
 
-
-        /// <summary>
-        ///     Register cancel callback. Two factors: either external cancellation token passed down from upper
-        ///     layer or timeout reached. Whichever comes first would trigger query cancellation.
-        /// </summary>
-        /// <param name="externalCancellationToken">cancellation token from upper layer</param>
-        internal void registerConnectionCancellationCallback(CancellationToken externalCancellationToken)
-        {
-            if (!externalCancellationToken.IsCancellationRequested)
-            {
-                externalCancellationToken.Register(() => { _connectionState = ConnectionState.Closed; });
-            }
-        }
-
         public bool IsStillRunning(QueryStatus status)
         {
             return QueryStatusExtensions.IsStillRunning(status);

Snowflake.Data/Client/SnowflakeDbConnectionPool.cs

@@ -16,7 +16,7 @@ public class SnowflakeDbConnectionPool
         private static IConnectionManager s_connectionManager;
         internal const ConnectionPoolType DefaultConnectionPoolType = ConnectionPoolType.MultipleConnectionPool;
 
-        private static IConnectionManager ConnectionManager
+        internal static IConnectionManager ConnectionManager
         {
             get
             {

Snowflake.Data/Core/SFStatement.cs

@@ -432,7 +432,12 @@ internal async Task<SFBaseResultSet> ExecuteAsync(int timeout, string sql, Dicti
 
                 return BuildResultSet(response, cancellationToken);
             }
-            catch
+            catch (OperationCanceledException ex)
+            {
+                logger.Warn($"Query execution canceled.");
+                throw;
+            }
+            catch (Exception ex)
             {
                 logger.Error("Query execution failed.");
                 throw;

Snowflake.Data/Core/Session/SessionPoolState.cs

@@ -19,6 +19,10 @@ public SessionPoolState(int idleSessionsCount, int busySessionsCount, int sessio
 
         public int Count() => _idleSessionsCount + _busySessionsCount + _sessionCreationsCount;
 
+        public int IdleSessionsCount { get => _idleSessionsCount; }
+
+        public int BusySessionsCount { get => _busySessionsCount; }
+
         public override string ToString()
         {
             return _extensiveFormat