From d0cca6046202ba15140c70b35dcf6f650dacd5f0 Mon Sep 17 00:00:00 2001
From: Shay Rojansky <roji@roji.org>
Date: Mon, 15 Nov 2021 21:16:44 +0100
Subject: [PATCH] Use read-write lock to protect concurrent pool accesses
 (#26658)

Fixes #26612
---
 .../SqliteConnectionFactory.cs                | 113 +++++++++++++++---
 1 file changed, 96 insertions(+), 17 deletions(-)

diff --git a/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs b/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs
index 43c10f98d18..75ce39a0800 100644
--- a/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs
+++ b/src/Microsoft.Data.Sqlite.Core/SqliteConnectionFactory.cs
@@ -11,17 +11,21 @@ internal class SqliteConnectionFactory
     {
         public static readonly SqliteConnectionFactory Instance = new();
 
+        private readonly bool _newLockingBehavior;
 #pragma warning disable IDE0052 // Remove unread private members
         private readonly Timer _pruneTimer;
 #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();
 
         protected SqliteConnectionFactory()
         {
-            if (!AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26422", out var enabled) || !enabled)
+            _newLockingBehavior = !AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26612", out var enabled) || !enabled;
+
+            if (!AppContext.TryGetSwitch("Microsoft.Data.Sqlite.Issue26422", out enabled) || !enabled)
             {
                 AppDomain.CurrentDomain.DomainUnload += (_, _) => ClearPools();
                 AppDomain.CurrentDomain.ProcessExit += (_, _) => ClearPools();
@@ -52,28 +56,63 @@ public SqliteConnectionInternal GetConnection(SqliteConnection outerConnection)
 
         public SqliteConnectionPoolGroup GetPoolGroup(string connectionString)
         {
-            if (!_poolGroups.TryGetValue(connectionString, out var poolGroup)
-                || (poolGroup.IsDisabled
-                    && !poolGroup.IsNonPooled))
+            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)
+                    {
+                        _lock.EnterWriteLock();
+                    }
+                    else
                     {
-                        var isNonPooled = connectionOptions.DataSource == ":memory:"
-                            || connectionOptions.Mode == SqliteOpenMode.Memory
-                            || connectionOptions.DataSource.Length == 0
-                            || !connectionOptions.Pooling;
+                        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 +132,33 @@ public void ReleasePool(SqliteConnectionPool pool, bool clearing)
 
         public void ClearPools()
         {
-            lock (this)
+            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 +188,16 @@ private void PruneCallback(object? _)
                 }
             }
 
-            lock (this)
+            if (_newLockingBehavior)
+            {
+                _lock.EnterWriteLock();
+            }
+            else
+            {
+                Monitor.Enter(this);
+            }
+
+            try
             {
                 var activePoolGroups = new Dictionary<string, SqliteConnectionPoolGroup>();
                 foreach (var entry in _poolGroups)
@@ -148,6 +216,17 @@ private void PruneCallback(object? _)
 
                 _poolGroups = activePoolGroups;
             }
+            finally
+            {
+                if (_newLockingBehavior)
+                {
+                    _lock.ExitWriteLock();
+                }
+                else
+                {
+                    Monitor.Exit(this);
+                }
+            }
         }
     }
 }