diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs index d155e67a2..b5f89b668 100644 --- a/Snowflake.Data/Core/Session/SessionPool.cs +++ b/Snowflake.Data/Core/Session/SessionPool.cs @@ -26,7 +26,9 @@ sealed class SessionPool : IDisposable private SecureString _password; private bool _pooling = true; private int _busySessions; - private bool _allowExceedMaxPoolSize = true; + private bool _allowExceedMaxPoolSize = true; // differentiates behavior when max size reached + private int _waitForIdleSessionSleepTimeout = 5; + private int _waitForIdleSessionTimeout = 120; internal SessionPool() { @@ -43,7 +45,7 @@ internal SessionPool(string connectionString, SecureString password) : this() { _connectionString = connectionString; _password = password; - _allowExceedMaxPoolSize = false; // TODO: SNOW-937190 + _allowExceedMaxPoolSize = false; } internal static SessionPool CreateSessionPoolV1() => new SessionPool(); @@ -133,6 +135,11 @@ private SFSession GetIdleSession(string connStr) private SFSession NewSession(String connectionString, SecureString password) { s_logger.Debug("SessionPool::NewSession"); + if (!_allowExceedMaxPoolSize && GetCurrentPoolSize() >= MaxPoolSize) + { + s_logger.Warn("Max pool size reached"); + return WaitForPooledSession(CancellationToken.None); + } try { var session = new SFSession(connectionString, password); @@ -156,6 +163,11 @@ private SFSession NewSession(String connectionString, SecureString password) private Task NewSessionAsync(String connectionString, SecureString password, CancellationToken cancellationToken) { s_logger.Debug("SessionPool::NewSessionAsync"); + if (!_allowExceedMaxPoolSize && GetCurrentPoolSize() >= MaxPoolSize) + { + s_logger.Warn("Max pool size reached"); + return Task.Run(() => WaitForPooledSession(cancellationToken), cancellationToken); + } var session = new SFSession(connectionString, password); _busySessions++; return session @@ -174,12 +186,46 @@ private Task NewSessionAsync(String connectionString, SecureString pa return session; }, TaskContinuationOptions.NotOnCanceled); } + + private SFSession WaitForPooledSession(CancellationToken cancellationToken) + { + if (!_pooling) + return null; + + long start = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + while (true) + { + if (cancellationToken.IsCancellationRequested) + cancellationToken.ThrowIfCancellationRequested(); + + lock (s_sessionPoolLock) + { + var session = GetIdleSession(_connectionString); + if (session != null) + { + _idleSessions.Remove(session); + _busySessions++; + return session; + } + } + + Thread.Sleep(_waitForIdleSessionSleepTimeout); + long now = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + if (start + _waitForIdleSessionTimeout < now) + throw new SnowflakeDbException( + new TimeoutException("No free connections in the pool."), // TODO: + SnowflakeDbException.CONNECTION_FAILURE_SSTATE, + SFError.REQUEST_TIMEOUT, + "Unable to connect."); + } + } internal bool AddSession(SFSession session) { s_logger.Debug("SessionPool::AddSession"); if (!_pooling) return false; + long timeNow = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); if (session.IsNotOpen() || session.IsExpired(_timeout, timeNow)) {