SNOW-937189 busy sessions counting in the pool

snowflakedb · Nov 13, 2023 · 186a6ea · 186a6ea
1 parent af36804
commit 186a6ea
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs
@@ -27,6 +27,7 @@ sealed class SessionPool : IDisposable
         internal string ConnectionString { get; }
         internal SecureString Password { get; }
         private bool _pooling = true;
+        private int _busySessions;
         private bool _allowExceedMaxPoolSize = true;
 
         private SessionPool()
@@ -36,6 +37,7 @@ private SessionPool()
                 _idleSessions = new List<SFSession>();
                 _maxPoolSize = MaxPoolSize;
                 _timeout = Timeout;
+                _busySessions = 0;
             }
         }
 
@@ -143,6 +145,7 @@ private SFSession NewSession(String connectionString, SecureString password)
             try
             {
                 var session = s_sessionFactory.NewSession(connectionString, password);
+                _busySessions++;
                 session.Open();
                 return session;
             }
@@ -163,6 +166,7 @@ private Task<SFSession> NewSessionAsync(String connectionString, SecureString pa
         {
             s_logger.Debug("SessionPool::NewSessionAsync");
             var session = s_sessionFactory.NewSession(connectionString, password);
+            _busySessions++;
             return session
                 .OpenAsync(cancellationToken)
                 .ContinueWith(previousTask =>
@@ -187,15 +191,19 @@ internal bool AddSession(SFSession session)
                 return false;
             long timeNow = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
             if (session.IsNotOpen() || session.IsExpired(_timeout, timeNow))
+            {
+                lock (s_sessionPoolLock)
+                {
+                    _busySessions = Math.Max(_busySessions - 1, 0);
+                }
                 return false;
+            }
 
             lock (s_sessionPoolLock)
             {
-                if (_idleSessions.Count >= _maxPoolSize)
-                {
-                    CleanExpiredSessions();
-                }
-                if (_idleSessions.Count >= _maxPoolSize)
+                _busySessions = Math.Max(_busySessions - 1, 0);
+                CleanExpiredSessions();
+                if (GetCurrentPoolSize() >= _maxPoolSize)
                 {
                     s_logger.Warn($"Pool is full - unable to add session with sid {session.sessionId}");
                     return false;
@@ -252,6 +260,8 @@ public long GetTimeout()
 
         public int GetCurrentPoolSize()
         {
+            if (!_allowExceedMaxPoolSize)
+                return _idleSessions.Count + _busySessions;
             return _idleSessions.Count;
         }