SNOW-902611 Fix busy session counting

snowflakedb · Sep 18, 2023 · dd608f8 · dd608f8
1 parent 64a40ed
commit dd608f8
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 12 deletions.
diff --git 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
@@ -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)