some improvements

snowflakedb · Jan 2, 2024 · 2bb889f · 2bb889f
1 parent b071399
commit 2bb889f
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 11 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/ConnectionMultiplePoolsIT.cs b/Snowflake.Data.Tests/IntegrationTests/ConnectionMultiplePoolsIT.cs
@@ -2,6 +2,8 @@
 using System.Data;
 using System.Diagnostics;
 using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
 using NUnit.Framework;
 using Snowflake.Data.Client;
 using Snowflake.Data.Core.Session;
@@ -102,7 +104,7 @@ public void TestReuseSessionInConnectionPoolReachingMaxConnections() // old name
         public void TestWaitForTheIdleConnectionWhenExceedingMaxConnectionsLimit()
         {
             // arrange
-            var connectionString = ConnectionString + "application=TestWaitForTheIdleConnectionWhenExceedingMaxConnectionsLimit";
+            var connectionString = ConnectionString;
             var pool = SnowflakeDbConnectionPool.GetPool(connectionString);
             Assert.AreEqual(0, pool.GetCurrentPoolSize(), "expecting pool to be empty");
             pool.SetMaxPoolSize(2);
@@ -126,6 +128,38 @@ public void TestWaitForTheIdleConnectionWhenExceedingMaxConnectionsLimit()
             conn1.Close();
             conn2.Close();
         }
+
+        [Test]
+        public void TestWaitForTheIdleConnectionWhenExceedingMaxConnectionsLimitAsync()
+        {
+            // arrange
+            var connectionString = ConnectionString;
+            var pool = SnowflakeDbConnectionPool.GetPool(connectionString);
+            Assert.AreEqual(0, pool.GetCurrentPoolSize(), "expecting pool to be empty");
+            pool.SetMaxPoolSize(2);
+            pool.SetWaitingTimeout(1000);
+            var conn1 = OpenedConnection(connectionString);
+            var conn2 = OpenedConnection(connectionString);
+            var watch = new Stopwatch();
+
+            // act
+            watch.Start();
+            var thrown = Assert.ThrowsAsync<SnowflakeDbException>(() => OpenedConnectionAsync(connectionString));
+            watch.Stop();
+
+            // assert
+            Assert.That(thrown.Message, Does.Contain("Unable to connect"));
+            Assert.IsTrue(thrown.InnerException is AggregateException);
+            var nextedException = ((AggregateException)thrown.InnerException).InnerException;
+            Assert.That(nextedException.Message, Does.Contain("Could not obtain a connection from the pool within a given timeout"));
+            Assert.GreaterOrEqual(watch.ElapsedMilliseconds, 1000);
+            Assert.LessOrEqual(watch.ElapsedMilliseconds, 1500);
+            Assert.AreEqual(pool.GetCurrentPoolSize(), 2);
+
+            // cleanup
+            conn1.Close();
+            conn2.Close();
+        }
 
         [Test]
         public void TestWaitInAQueueForAnIdleSession()
@@ -266,5 +300,13 @@ private SnowflakeDbConnection OpenedConnection(string connectionString)
             connection.Open();
             return connection;
         }
+
+        private async Task<SnowflakeDbConnection> OpenedConnectionAsync(string connectionString)
+        {
+            var connection = new SnowflakeDbConnection();
+            connection.ConnectionString = connectionString;
+            await connection.OpenAsync().ConfigureAwait(false);
+            return connection;
+        }
     }
 }
diff --git a/Snowflake.Data/Core/Session/CreateSessionTokens.cs b/Snowflake.Data/Core/Session/CreateSessionTokens.cs
@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using Snowflake.Data.Core.Session;
 
 namespace Snowflake.Data.Core.Session
 {
@@ -40,4 +39,4 @@ public int Count()
             return _tokenCount;
         }
     }
-}
+}
diff --git a/Snowflake.Data/Core/Session/NotCountingCreateSessionTokens.cs b/Snowflake.Data/Core/Session/NotCountingCreateSessionTokens.cs
@@ -2,7 +2,9 @@ namespace Snowflake.Data.Core.Session
 {
     internal class NotCountingCreateSessionTokens: ICreateSessionTokens
     {
-        public CreateSessionToken BeginCreate() => new CreateSessionToken(0);
+        private const int IrrelevantCreateSessionTimeout = 0;
+
+        public CreateSessionToken BeginCreate() => new CreateSessionToken(IrrelevantCreateSessionTimeout);
 
         public void EndCreate(CreateSessionToken token)
         {

diff --git a/Snowflake.Data/Core/Session/SemaphoreBasedQueue.cs b/Snowflake.Data/Core/Session/SemaphoreBasedQueue.cs
@@ -20,7 +20,7 @@ public bool Wait(int millisecondsTimeout, CancellationToken cancellationToken)
             {
                 return _semaphore.Wait(millisecondsTimeout, cancellationToken);
             }
-            catch (OperationCanceledException exception)
+            catch (OperationCanceledException)
             {
                 return false;
             }

diff --git a/Snowflake.Data/Core/Session/SessionPool.cs b/Snowflake.Data/Core/Session/SessionPool.cs
@@ -22,6 +22,7 @@ sealed class SessionPool : IDisposable
         private readonly List<SFSession> _idleSessions;
         private readonly IWaitingQueue _waitingQueue;
         private readonly ICreateSessionTokens _createSessionTokens;
+        private readonly ICreateSessionTokens _noPoolingCreateSessionTokens = new NotCountingCreateSessionTokens();
         private int _maxPoolSize;
         private long _timeout;
         private const int MaxPoolSize = 10;
@@ -100,20 +101,18 @@ internal SFSession GetSession(string connStr, SecureString password)
         {
             s_logger.Debug("SessionPool::GetSession");
             if (!_pooling)
-                return NewSession(connStr, password, new CreateSessionToken(0));
+                return NewSession(connStr, password, _noPoolingCreateSessionTokens.BeginCreate());
             var sessionOrCreateToken = GetIdleSession(connStr);
             return sessionOrCreateToken.Session ?? NewSession(connStr, password, sessionOrCreateToken.CreateToken);
         }
 
-        internal Task<SFSession> GetSessionAsync(string connStr, SecureString password, CancellationToken cancellationToken)
+        internal async Task<SFSession> GetSessionAsync(string connStr, SecureString password, CancellationToken cancellationToken)
         {
             s_logger.Debug("SessionPool::GetSessionAsync");
             if (!_pooling)
-                return NewSessionAsync(connStr, password, new CreateSessionToken(0), cancellationToken);
+                return await NewSessionAsync(connStr, password, _noPoolingCreateSessionTokens.BeginCreate(), cancellationToken).ConfigureAwait(false);
             var sessionOrCreateToken = GetIdleSession(connStr);
-            return sessionOrCreateToken.Session != null
-                ? Task.FromResult(sessionOrCreateToken.Session)
-                : NewSessionAsync(connStr, password, sessionOrCreateToken.CreateToken, cancellationToken);
+            return sessionOrCreateToken.Session ?? await NewSessionAsync(connStr, password, sessionOrCreateToken.CreateToken, cancellationToken).ConfigureAwait(false);
         }
 
         internal SFSession GetSession() => GetSession(ConnectionString, Password);