SNOW-937189 busy sessions counting in the pool

snowflakedb · Oct 11, 2023 · 04914ff · 04914ff
1 parent 73bb45c
commit 04914ff
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
@@ -25,6 +25,7 @@ sealed class SessionPool : IDisposable
         private string _connectionString;
         private SecureString _password;
         private bool _pooling = true;
+        private int _busySessions;
         private bool _allowExceedMaxPoolSize = true;
 
         internal SessionPool()
@@ -34,6 +35,7 @@ internal SessionPool()
                 _idleSessions = new List<SFSession>();
                 _maxPoolSize = MaxPoolSize;
                 _timeout = Timeout;
+                _busySessions = 0;
             }
         }
 
@@ -134,6 +136,7 @@ private SFSession NewSession(String connectionString, SecureString password)
             try
             {
                 var session = new SFSession(connectionString, password);
+                _busySessions++;
                 session.Open();
                 return session;
             }
@@ -154,6 +157,7 @@ private Task<SFSession> NewSessionAsync(String connectionString, SecureString pa
         {
             s_logger.Debug("SessionPool::NewSessionAsync");
             var session = new SFSession(connectionString, password);
+            _busySessions++;
             return session
                 .OpenAsync(cancellationToken)
                 .ContinueWith(previousTask =>
@@ -178,15 +182,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;
@@ -233,6 +241,8 @@ public long GetTimeout()
 
         public int GetCurrentPoolSize()
         {
+            if (!_allowExceedMaxPoolSize)
+                return _idleSessions.Count + _busySessions;
             return _idleSessions.Count;
         }