fix cancellation case

snowflakedb · Jan 26, 2024 · 284b36a · 284b36a
1 parent 9f585a2
commit 284b36a
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs
@@ -293,6 +293,10 @@ private Task<SFSession> NewSessionAsync(String connectionString, SecureString pa
             return session
                 .OpenAsync(cancellationToken)
                 .ContinueWith(previousTask =>
+                {
+                    _sessionCreationTokenCounter.RemoveToken(sessionCreationToken);
+                }, TaskContinuationOptions.OnlyOnCanceled)
+                .ContinueWith(previousTask =>
                 {
                     if (previousTask.IsFaulted)
                     {
@@ -319,7 +323,7 @@ private Task<SFSession> NewSessionAsync(String connectionString, SecureString pa
                     _sessionPoolEventHandler.OnNewSessionCreated(this);
                     _sessionPoolEventHandler.OnSessionProvided(this);
                     return session;
-                }, TaskContinuationOptions.None); // previously it was NotOnCanceled but we would like to execute it even in case of cancellation to properly update counters 
+                }, TaskContinuationOptions.NotOnCanceled); 
         }
 
         internal bool AddSession(SFSession session)