Use read-write lock to protect concurrent pool accesses

Fixes dotnet#26612
roji · Nov 12, 2021 · c5b0c1a · c5b0c1a
1 parent b38c269
commit c5b0c1a
Showing 1 changed file with 98 additions and 16 deletions.
diff --git a/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs b/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs
@@ -16,6 +16,7 @@ internal class SqliteConnectionFactory
 #pragma warning restore IDE0052 // Remove unread private members
         private readonly List<SqliteConnectionPoolGroup> _idlePoolGroups = new();
         private readonly List<SqliteConnectionPool> _poolsToRelease = new();
+        private readonly ReaderWriterLockSlim _lock = new();
 
         private Dictionary<string, SqliteConnectionPoolGroup> _poolGroups = new();
 
@@ -52,28 +53,65 @@ public SqliteConnectionInternal GetConnection(SqliteConnection outerConnection)
 
         public SqliteConnectionPoolGroup GetPoolGroup(string connectionString)
         {
-            if (!_poolGroups.TryGetValue(connectionString, out var poolGroup)
-                || (poolGroup.IsDisabled
-                    && !poolGroup.IsNonPooled))
+            var newLockingBehavior = !AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26612", out var enabled) || !enabled;
+
+            if (newLockingBehavior)
             {
-                var connectionOptions = new SqliteConnectionStringBuilder(connectionString);
+                _lock.EnterUpgradeableReadLock();
+            }
 
-                lock (this)
+            try
+            {
+                if (!_poolGroups.TryGetValue(connectionString, out var poolGroup)
+                    || (poolGroup.IsDisabled
+                        && !poolGroup.IsNonPooled))
                 {
-                    if (!_poolGroups.TryGetValue(connectionString, out poolGroup))
+                    var connectionOptions = new SqliteConnectionStringBuilder(connectionString);
+
+                    if (newLockingBehavior)
                     {
-                        var isNonPooled = connectionOptions.DataSource == ":memory:"
-                            || connectionOptions.Mode == SqliteOpenMode.Memory
-                            || connectionOptions.DataSource.Length == 0
-                            || !connectionOptions.Pooling;
+                        _lock.EnterWriteLock();
+                    }
+                    else
+                    {
+                        Monitor.Enter(this);
+                    }
 
-                        poolGroup = new SqliteConnectionPoolGroup(connectionOptions, connectionString, isNonPooled);
-                        _poolGroups.Add(connectionString, poolGroup);
+                    try
+                    {
+                        if (!_poolGroups.TryGetValue(connectionString, out poolGroup))
+                        {
+                            var isNonPooled = connectionOptions.DataSource == ":memory:"
+                                || connectionOptions.Mode == SqliteOpenMode.Memory
+                                || connectionOptions.DataSource.Length == 0
+                                || !connectionOptions.Pooling;
+
+                            poolGroup = new SqliteConnectionPoolGroup(connectionOptions, connectionString, isNonPooled);
+                            _poolGroups.Add(connectionString, poolGroup);
+                        }
+                    }
+                    finally
+                    {
+                        if (newLockingBehavior)
+                        {
+                            _lock.ExitWriteLock();
+                        }
+                        else
+                        {
+                            Monitor.Exit(this);
+                        }
                     }
                 }
-            }
 
-            return poolGroup;
+                return poolGroup;
+            }
+            finally
+            {
+                if (newLockingBehavior)
+                {
+                    _lock.ExitUpgradeableReadLock();
+                }
+            }
         }
 
         public void ReleasePool(SqliteConnectionPool pool, bool clearing)
@@ -93,13 +131,35 @@ public void ReleasePool(SqliteConnectionPool pool, bool clearing)
 
         public void ClearPools()
         {
-            lock (this)
+            var newLockingBehavior = !AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26612", out var enabled) || !enabled;
+
+            if (newLockingBehavior)
+            {
+                _lock.EnterWriteLock();
+            }
+            else
+            {
+                Monitor.Enter(this);
+            }
+
+            try
             {
                 foreach (var entry in _poolGroups)
                 {
                     entry.Value.Clear();
                 }
             }
+            finally
+            {
+                if (newLockingBehavior)
+                {
+                    _lock.ExitWriteLock();
+                }
+                else
+                {
+                    Monitor.Exit(this);
+                }
+            }
         }
 
         private void PruneCallback(object? _)
@@ -129,7 +189,18 @@ private void PruneCallback(object? _)
                 }
             }
 
-            lock (this)
+            var newLockingBehavior = !AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26612", out var enabled) || !enabled;
+
+            if (newLockingBehavior)
+            {
+                _lock.EnterWriteLock();
+            }
+            else
+            {
+                Monitor.Enter(this);
+            }
+
+            try
             {
                 var activePoolGroups = new Dictionary<string, SqliteConnectionPoolGroup>();
                 foreach (var entry in _poolGroups)
@@ -148,6 +219,17 @@ private void PruneCallback(object? _)
 
                 _poolGroups = activePoolGroups;
             }
+            finally
+            {
+                if (newLockingBehavior)
+                {
+                    _lock.ExitWriteLock();
+                }
+                else
+                {
+                    Monitor.Exit(this);
+                }
+            }
         }
     }
 }