From dd608f8fce85ccf1f163ce3ba003673b16e3b8d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Hofman?= <michal.hofman@snowflake.com>
Date: Mon, 18 Sep 2023 17:14:43 +0200
Subject: [PATCH] SNOW-902611 Fix busy session counting

---
 .../Client/SnowflakeDbConnectionPool.cs       |  7 +++-
 Snowflake.Data/Core/Session/SessionPool.cs    | 41 ++++++++++++++-----
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/Snowflake.Data/Client/SnowflakeDbConnectionPool.cs b/Snowflake.Data/Client/SnowflakeDbConnectionPool.cs
index 11241c7c2..e253962ac 100644
--- a/Snowflake.Data/Client/SnowflakeDbConnectionPool.cs
+++ b/Snowflake.Data/Client/SnowflakeDbConnectionPool.cs
@@ -108,7 +108,10 @@ internal static bool AddSession(string connectionString, SecureString password,
             s_logger.Debug("SnowflakeDbConnectionPool::AddSession");
             return Instance.AddSession(connectionString, password, session);
         }
-
-
+        
+        internal static PoolManagerVersion GetVersion()
+        {
+            return s_poolVersion;
+        }
     }
 }
diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs
index 6fe2fa757..ac54b5553 100644
--- a/Snowflake.Data/Core/Session/SessionPool.cs
+++ b/Snowflake.Data/Core/Session/SessionPool.cs
@@ -36,6 +36,7 @@ internal SessionPool(string connectionString, SecureString password)
             lock (s_sessionPoolLock)
             {
                 _sessionPool = new List<SFSession>();
+                _busySessions = 0;
                 _maxPoolSize = MaxPoolSize;
                 _timeout = Timeout;
                 // _minPoolSize = MinPoolSize; // TODO:
@@ -321,23 +322,36 @@ internal bool AddSession(SFSession session)
                 return false;
             long timeNow = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
             if (session.IsNotOpen() || session.IsExpired(_timeout, timeNow))
-                return false; // TODO: fix because it is counted in the pool 
+            {
+                s_logger.Warn($"Session returning to the pool in an undesired state: {session.sessionId}");
+                // TODO: fix because it is counted in the pool 
+                // TODO: lock
+                if (_busySessions > 0)
+                _busySessions--; 
+                return false; 
+            }
 
+            if (_sessionPool.Count >= _maxPoolSize)
+                CleanExpiredSessions();
+            
             lock (s_sessionPoolLock)
             {
                 if (_sessionPool.Count >= _maxPoolSize)
                 {
-                    CleanExpiredSessions();
-                }
-                if (_sessionPool.Count >= _maxPoolSize)
-                {
-                    // pool is full
+                    s_logger.Warn($"Pool is full, cannot add session with sid {session.sessionId}"); 
                     return false;
                 }
 
-                s_logger.Debug($"pool connection with sid {session.sessionId}");
-                _sessionPool.Add(session);
-                return true;
+                if (_busySessions > 0)
+                {
+                    _busySessions--;
+                    s_logger.Debug($"Connection returned to the pool with sid {session.sessionId}");
+                    _sessionPool.Add(session);
+                    return true;
+                }
+                
+                s_logger.Warn($"Unexpected session with sid {session.sessionId} was not returned to the pool"); // or clear pool was called and session was created before
+                return false;
             }
         }
 
@@ -351,6 +365,7 @@ internal void ClearAllPools()
                     session.close();
                 }
                 _sessionPool.Clear();
+                _busySessions = 0; // TODO: check test TestConnectionPoolIsFull
             }
         }
 
@@ -376,7 +391,13 @@ public long GetTimeout()
 
         public int GetCurrentPoolSize()
         {
-            return _sessionPool.Count + _busySessions;
+            switch (SnowflakeDbConnectionPool.GetVersion())
+            {
+                case PoolManagerVersion.Version1: return _sessionPool.Count;
+                case PoolManagerVersion.Version2: return _sessionPool.Count + _busySessions;
+            }
+            throw new NotSupportedException("Unknown pool version");
+
         }
 
         public bool SetPooling(bool isEnable)