From 284b36a2bd6e3b477de9660a6304d7d229744114 Mon Sep 17 00:00:00 2001
From: Krzysztof Nozderko <krzysztof.nozderko@snowflake.com>
Date: Fri, 26 Jan 2024 16:17:21 +0100
Subject: [PATCH] fix cancellation case

---
 Snowflake.Data/Core/Session/SessionPool.cs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs
index b1fd4393b..8bd13284b 100644
--- 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)