Merge pull request #16 from madelson/fix-azure-connection-leak

Fix azure connection leak
madelson · Feb 3, 2018 · 46d11e6 · 46d11e6
2 parents 1e0c670 + dbde988
commit 46d11e6
Show file tree

Hide file tree

Showing 8 changed files with 188 additions and 16 deletions.
diff --git a/DistributedLock.Tests/AbstractTestCases/ConnectionStringStrategyTestCases.cs b/DistributedLock.Tests/AbstractTestCases/ConnectionStringStrategyTestCases.cs
@@ -0,0 +1,80 @@
+using Medallion.Threading.Tests.Sql;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using System;
+using System.Collections.Generic;
+using System.Data.SqlClient;
+using System.Diagnostics;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Medallion.Threading.Tests
+{
+    [TestClass]
+    public abstract class ConnectionStringStrategyTestCases<TEngineFactory, TConnectionManagementProvider> : TestBase
+        where TEngineFactory : ITestingSqlDistributedLockEngineFactory, new()
+        where TConnectionManagementProvider : ConnectionStringProvider, new()
+    {
+        /// <summary>
+        /// Tests that internally-owned connections are properly cleaned up by disposing the lock handle 
+        /// </summary>
+        [TestMethod]
+        public void TestConnectionDoesNotLeak()
+        {
+            var applicationName = nameof(TestConnectionDoesNotLeak) + Guid.NewGuid();
+            var connectionString = new SqlConnectionStringBuilder(ConnectionStringProvider.ConnectionString)
+                {
+                    ApplicationName = applicationName,
+                }
+                .ConnectionString;
+
+            using (ConnectionStringProvider.UseConnectionString(connectionString))
+            using (var engine = this.CreateEngine())
+            {
+                var @lock = engine.CreateLock(nameof(TestConnectionDoesNotLeak));
+
+                for (var i = 0; i < 30; ++i)
+                {
+                    using (@lock.Acquire())
+                    {
+                        CountActiveSessions().ShouldEqual(1, this.GetType().Name);
+                    }
+                    // still alive due to pooling
+                    CountActiveSessions().ShouldEqual(1, this.GetType().Name);
+                }
+            }
+
+            using (var connection = new SqlConnection(connectionString))
+            {
+                SqlConnection.ClearPool(connection);
+                // checking immediately seems flaky; likely clear pool finishing
+                // doesn't guarantee that SQL will immediately reflect the clear
+                var maxWaitForPoolsToClear = TimeSpan.FromSeconds(5);
+                var stopwatch = Stopwatch.StartNew();
+                do
+                {
+                    var activeCount = CountActiveSessions();
+                    if (activeCount == 0) { return; }
+                    Thread.Sleep(25);
+                }
+                while (stopwatch.Elapsed < maxWaitForPoolsToClear);
+            }
+
+            int CountActiveSessions()
+            {
+                using (var connection = new SqlConnection(ConnectionStringProvider.ConnectionString))
+                {
+                    connection.Open();
+                    using (var command = connection.CreateCommand())
+                    {
+                        command.CommandText = $@"SELECT COUNT(*) FROM sys.dm_exec_sessions WHERE program_name = '{applicationName}'";
+                        return (int)command.ExecuteScalar();
+                    }
+                }
+            }
+        }
+
+        private TestingDistributedLockEngine CreateEngine() => new TEngineFactory().Create<TConnectionManagementProvider>();
+    }
+}
diff --git a/DistributedLock.Tests/DistributedLock.Tests.csproj b/DistributedLock.Tests/DistributedLock.Tests.csproj
@@ -36,12 +36,18 @@
     <WarningLevel>4</WarningLevel>
   </PropertyGroup>
   <ItemGroup>
+    <Reference Include="MedallionCollections, Version=1.0.0.0, Culture=neutral, processorArchitecture=MSIL">
+      <HintPath>..\packages\MedallionCollections.1.1.0\lib\net45\MedallionCollections.dll</HintPath>
+    </Reference>
     <Reference Include="MedallionShell, Version=1.2.1.0, Culture=neutral, processorArchitecture=MSIL">
       <HintPath>..\packages\MedallionShell.1.2.1\lib\net45\MedallionShell.dll</HintPath>
       <Private>True</Private>
     </Reference>
     <Reference Include="System" />
     <Reference Include="System.Data" />
+    <Reference Include="System.ValueTuple, Version=4.0.2.0, Culture=neutral, PublicKeyToken=cc7b13ffcd2ddd51, processorArchitecture=MSIL">
+      <HintPath>..\packages\System.ValueTuple.4.4.0\lib\netstandard1.0\System.ValueTuple.dll</HintPath>
+    </Reference>
   </ItemGroup>
   <Choose>
     <When Condition="('$(VisualStudioVersion)' == '10.0' or '$(VisualStudioVersion)' == '') and '$(TargetFrameworkVersion)' == 'v3.5'">
@@ -56,6 +62,7 @@
     </Otherwise>
   </Choose>
   <ItemGroup>
+    <Compile Include="AbstractTestCases\ConnectionStringStrategyTestCases.cs" />
     <Compile Include="AbstractTestCases\SqlDistributedSemaphoreSelfDeadlockTestCases.cs" />
     <Compile Include="Infrastructure\ActionRegistrationDisposable.cs" />
     <Compile Include="AbstractTestCases\DistributedLockCoreTestCases.cs" />

diff --git a/DistributedLock.Tests/Infrastructure/TestingSqlConnectionManagementProviders.cs b/DistributedLock.Tests/Infrastructure/TestingSqlConnectionManagementProviders.cs
@@ -23,20 +23,28 @@ public abstract class ConnectionStringProvider : TestingSqlConnectionManagementP
             }
             .ConnectionString;
 
+        private static volatile string currentConnectionString = ConnectionString;
+
         private readonly SqlDistributedLockConnectionStrategy? _strategy;
 
         public ConnectionStringProvider(SqlDistributedLockConnectionStrategy? strategy)
         {
             this._strategy = strategy;
         }
 
-        public override ConnectionInfo GetConnectionInfo() => new ConnectionInfo { ConnectionString = ConnectionString, Strategy = this._strategy };
+        public override ConnectionInfo GetConnectionInfo() => new ConnectionInfo { ConnectionString = currentConnectionString, Strategy = this._strategy };
 
         /// <summary>
         /// Since every lock handle has a dedicated connection (even in the case of multiplexing we don't share a connection for two acquires
         /// on the same name, we are not reentrant)
         /// </summary>
         internal override bool IsReentrantForAppLock => false;
+
+        public static IDisposable UseConnectionString(string connectionString)
+        {
+            currentConnectionString = connectionString;
+            return new ReleaseAction(() => currentConnectionString = ConnectionString);
+        }
     }
 
     public sealed class NoStrategyConnectionStringProvider : ConnectionStringProvider

diff --git a/DistributedLock.Tests/Tests/CombinatorialTests.cs b/DistributedLock.Tests/Tests/CombinatorialTests.cs
@@ -186,4 +186,58 @@ public class SqlSemaphoreSelfDeadlock_ConnectionProviderTest : SqlDistributedSem
 
     [TestClass]
     public class SqlSemaphoreSelfDeadlock_TransactionProviderTest : SqlDistributedSemaphoreSelfDeadlockTestCases<TransactionProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_NoStrategyConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, NoStrategyConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_DefaultConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, DefaultConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_AzureConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, AzureConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_ConnectionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, ConnectionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_TransactionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, TransactionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlEngineFactory_MultiplexedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedLockEngineFactory, MultiplexedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_NoStrategyConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, NoStrategyConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_DefaultConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, DefaultConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_AzureConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, AzureConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_ConnectionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, ConnectionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_TransactionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, TransactionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlReaderWriterEngineFactory_MultiplexedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedReaderWriterLockEngineFactory, MultiplexedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_NoStrategyConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, NoStrategyConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_DefaultConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, DefaultConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_AzureConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, AzureConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_ConnectionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, ConnectionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_TransactionBasedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, TransactionBasedConnectionStringProvider> { }
+
+    [TestClass]
+    public class ConnectionStringStrategy_SqlSemaphoreEngineFactory_MultiplexedConnectionStringProviderTest : ConnectionStringStrategyTestCases<TestingSqlDistributedSemaphoreEngineFactory, MultiplexedConnectionStringProvider> { }
 }
diff --git a/DistributedLock.Tests/Tests/TestSetupTest.cs b/DistributedLock.Tests/Tests/TestSetupTest.cs
@@ -1,4 +1,5 @@
-using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Medallion.Collections;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -69,18 +70,31 @@ public class {testClassName} : {GetCSharpName(testClassType)} {{ }}";
 
         private static string RemoveGenericMarkers(string name) => Regex.Replace(name, @"`\d+", string.Empty);
 
-        private List<Type> GetTypesForGenericParameters(Type[] genericParameters)
+        private List<Type[]> GetTypesForGenericParameters(Type[] genericParameters)
         {
-            if (genericParameters.Length > 1) { throw new NotSupportedException(); }
+            var genericParameterTypes = genericParameters.Select(this.GetTypesForGenericParameter).ToArray();
+            var allCombinations = Traverse.DepthFirst(
+                    root: (index: 0, value: Enumerable.Empty<Type>()),
+                    children: t => t.index == genericParameterTypes.Length
+                        ? Enumerable.Empty<(int index, IEnumerable<Type> value)>()
+                        : genericParameterTypes[t.index].Select(type => (index: t.index + 1, value: t.value.Append(type)))
+                )
+                .Where(t => t.index == genericParameterTypes.Length)
+                .Select(t => t.value.ToArray())
+                .ToList();
+            return allCombinations;
+        }
 
-            var constraints = genericParameters[0].GetGenericParameterConstraints();
+        private Type[] GetTypesForGenericParameter(Type genericParameter)
+        {
+            var constraints = genericParameter.GetGenericParameterConstraints();
             return this.GetType().Assembly
                 .GetTypes()
                 // this doesn't support fancy constraints like class
                 // see https://stackoverflow.com/questions/4864496/checking-if-an-object-meets-a-generic-parameter-constraint
                 .Where(t => !t.IsAbstract && constraints.All(c => c.IsAssignableFrom(t)))
                 .SelectMany(t => t.IsGenericTypeDefinition ? this.GetTypesForGenericParameters(t.GetGenericArguments()).Select(p => t.MakeGenericType(p)) : new[] { t })
-                .ToList();
+                .ToArray();
         }
     }
 }
diff --git a/DistributedLock.Tests/packages.config b/DistributedLock.Tests/packages.config
@@ -1,4 +1,6 @@
 <?xml version="1.0" encoding="utf-8"?>
 <packages>
+  <package id="MedallionCollections" version="1.1.0" targetFramework="net46" />
   <package id="MedallionShell" version="1.2.1" targetFramework="net45" />
+  <package id="System.ValueTuple" version="4.4.0" targetFramework="net46" />
 </packages>
diff --git a/DistributedLock/Sql/InternalLocks/AzureSqlDistributedLock.cs b/DistributedLock/Sql/InternalLocks/AzureSqlDistributedLock.cs
@@ -31,7 +31,7 @@ public IDisposable TryAcquire<TLockCookie>(int timeoutMillis, ISqlSynchronizatio
                 {
                     var internalHandle = lockScope.InternalLock.TryAcquire(timeoutMillis, strategy, contextHandle: lockScope.InternalHandle);
                     return internalHandle != null
-                        ? new LockScope(internalHandle, lockScope.InternalLock, lockScope.Keepalive, ownsKeepalive: false)
+                        ? new LockScope(internalHandle, lockScope.InternalLock, lockScope.Keepalive, connection: null)
                         : null;
                 }
                 finally
@@ -52,7 +52,7 @@ public IDisposable TryAcquire<TLockCookie>(int timeoutMillis, ISqlSynchronizatio
                 {
                     var keepalive = new KeepaliveHelper(connection);
                     keepalive.Start();
-                    result = new LockScope(internalHandle, internalLock, keepalive, ownsKeepalive: true);
+                    result = new LockScope(internalHandle, internalLock, keepalive, connection);
                 }
             }
             finally
@@ -78,7 +78,7 @@ public async Task<IDisposable> TryAcquireAsync<TLockCookie>(int timeoutMillis, I
                 {
                     var internalHandle = await lockScope.InternalLock.TryAcquireAsync(timeoutMillis, strategy, cancellationToken, contextHandle: lockScope.InternalHandle).ConfigureAwait(false);
                     return internalHandle != null
-                        ? new LockScope(internalHandle, lockScope.InternalLock, lockScope.Keepalive, ownsKeepalive: false)
+                        ? new LockScope(internalHandle, lockScope.InternalLock, lockScope.Keepalive, connection: null)
                         : null;
                 }
                 finally
@@ -99,7 +99,7 @@ public async Task<IDisposable> TryAcquireAsync<TLockCookie>(int timeoutMillis, I
                 {
                     var keepalive = new KeepaliveHelper(connection);
                     keepalive.Start();
-                    result = new LockScope(internalHandle, internalLock, keepalive, ownsKeepalive: true);
+                    result = new LockScope(internalHandle, internalLock, keepalive, connection);
                 }
             }
             finally
@@ -109,22 +109,22 @@ public async Task<IDisposable> TryAcquireAsync<TLockCookie>(int timeoutMillis, I
 
             return result;
         }
-
+        
         private sealed class LockScope : IDisposable
         {
-            private readonly bool ownsKeepalive;
             private KeepaliveHelper keepalive;
+            private SqlConnection connection;
 
             public LockScope(
                 IDisposable internalHandle, 
                 ExternalConnectionOrTransactionSqlDistributedLock internalLock,
                 KeepaliveHelper keepalive,
-                bool ownsKeepalive)
+                SqlConnection connection)
             {
                 this.InternalHandle = internalHandle;
                 this.InternalLock = internalLock;
                 this.keepalive = keepalive;
-                this.ownsKeepalive = ownsKeepalive;
+                this.connection = connection;
             }
 
             public IDisposable InternalHandle { get; private set; }
@@ -144,9 +144,15 @@ public void Dispose()
                         this.InternalHandle.Dispose();
                         this.InternalHandle = null;
                         this.InternalLock = null;
-                        if (!this.ownsKeepalive)
+                        if (this.connection != null)
+                        {
+                            // if we own the connection then dispose it
+                            this.connection.Dispose();
+                            this.connection = null;
+                        }
+                        else
                         {
-                            // if the keepalive is owned by an outer handle, restart it
+                            // otherwise, the keepalive is owned by an outer handle so restart it
                             keepalive.Start();
                         }
                     }

diff --git a/README.md b/README.md
@@ -199,6 +199,7 @@ Most of the time, you'll want to use the default connection strategy. See more d
 ## Release notes
 - 1.4.0
 	- Added a SQL-based distributed semaphore ([#7](https://github.com/madelson/DistributedLock/issues/7))
+	- Fix bug where SqlDistributedLockConnectionStrategy.Azure would leak connections, relying on GC to reclaim them ([#14](https://github.com/madelson/DistributedLock/issues/14)). Thanks [zavalita1](https://github.com/zavalita1) for investigating this issue!
 - 1.3.1 Minor fix to avoid "leaking" isolation level changes in transaction-based locks ([#8](https://github.com/madelson/DistributedLock/issues/8)). Also switched to the VS2017 project file format
 - 1.3.0 Added an Azure connection strategy to keep lock connections from becoming idle and being reclaimed by Azure's connection governor ([#5](https://github.com/madelson/DistributedLock/issues/5))
 - 1.2.0