Simplify

src/Workspaces/Core/Portable/Storage/AbstractPersistentStorageService.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Concurrent;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,11 +22,8 @@ internal abstract partial class AbstractPersistentStorageService(IPersistentStor
 {
     protected readonly IPersistentStorageConfiguration Configuration = configuration;
 
-    /// <summary>
-    /// This lock guards all mutable fields in this type.
-    /// </summary>
     private readonly SemaphoreSlim _lock = new(initialCount: 1);
-    private IChecksummedPersistentStorage? _currentPersistentStorage;
+    private readonly ConcurrentDictionary<SolutionKey, IChecksummedPersistentStorage> _solutionKeyToStorage = new();
 
     protected abstract string GetDatabaseFilePath(string workingFolderPath);
 
@@ -42,11 +40,9 @@ public async ValueTask<IChecksummedPersistentStorage> GetStorageAsync(SolutionKe
         if (solutionKey.FilePath == null)
             return NoOpPersistentStorage.GetOrThrow(solutionKey, Configuration.ThrowOnFailure);
 
-        // Without taking the lock, see if we can use the last storage system we were asked to create.  Ensure we use a
-        // using so that if we don't take it we still release this reference count.
-        var existing = _currentPersistentStorage;
-        if (solutionKey == existing?.SolutionKey)
-            return existing;
+        // Without taking the lock, see if we can lookup a storage for this key.
+        if (_solutionKeyToStorage.TryGetValue(solutionKey, out var storage))
+            return storage;
 
         var workingFolder = Configuration.TryGetStorageLocation(solutionKey);
         if (workingFolder == null)
@@ -55,12 +51,13 @@ public async ValueTask<IChecksummedPersistentStorage> GetStorageAsync(SolutionKe
         using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
         {
             // See if another thread set to the solution we care about while we were waiting on the lock.
-            if (solutionKey != _currentPersistentStorage?.SolutionKey)
+            if (!_solutionKeyToStorage.TryGetValue(solutionKey, out storage))
             {
-                _currentPersistentStorage = await CreatePersistentStorageAsync(solutionKey, workingFolder, cancellationToken).ConfigureAwait(false);
+                storage = await CreatePersistentStorageAsync(solutionKey, workingFolder, cancellationToken).ConfigureAwait(false);
+                _solutionKeyToStorage.Add(solutionKey, storage);
             }
 
-            return _currentPersistentStorage;
+            return storage;
         }
     }
 
@@ -120,8 +117,7 @@ private void Shutdown(CancellationToken cancellationToken)
     {
         using (_lock.DisposableWait(cancellationToken))
         {
-            // We will transfer ownership in a thread-safe way out so we can dispose outside the lock
-            _currentPersistentStorage = null;
+            _solutionKeyToStorage.Clear();
         }
     }
 

src/Workspaces/Core/Portable/Storage/SQLite/v2/Interop/SqlConnection.cs

@@ -22,15 +22,15 @@ namespace Microsoft.CodeAnalysis.SQLite.v2.Interop;
 /// to open the DB if it exists, or create it if it does not.
 /// 
 /// Connections are considered relatively heavyweight and are pooled (see <see
-/// cref="SQLiteConnectionPool.GetPooledConnection(out SqlConnection)"/>).  Connections can be used by different
+/// cref="SQLitePersistentStorage.GetPooledConnection(out SqlConnection)"/>).  Connections can be used by different
 /// threads, but only as long as they are used by one thread at a time.  They are not safe for concurrent use by several
 /// threads.
 /// 
 /// <see cref="SqlStatement"/>s can be created through the user of <see cref="GetResettableStatement"/>.
 /// These statements are cached for the lifetime of the connection and are only finalized
 /// (i.e. destroyed) when the connection is closed.
 /// </summary>
-internal class SqlConnection
+internal sealed class SqlConnection
 {
     // Cached UTF-8 (and null terminated) versions of the common strings we need to pass to sqlite.  Used to prevent
     // having to convert these names to/from utf16 to UTF-8 on every call.  Sqlite requires these be null terminated.

src/Workspaces/Core/Portable/Storage/SQLite/v2/SQLiteConnectionPool.cs → src/Workspaces/Core/Portable/Storage/SQLite/v2/SQLitePersistentStorage+PooledConnection.cs

@@ -3,38 +3,20 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Generic;
-using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.SQLite.v2.Interop;
 
 namespace Microsoft.CodeAnalysis.SQLite.v2;
 
-internal sealed partial class SQLiteConnectionPool(
-    SQLiteConnectionPoolService connectionPoolService, IPersistentStorageFaultInjector? faultInjector, string databasePath, object ownershipLock)
+internal partial class SQLitePersistentStorage
 {
-    // We pool connections to the DB so that we don't have to take the hit of 
-    // reconnecting.  The connections also cache the prepared statements used
-    // to get/set data from the db.  A connection is safe to use by one thread
-    // at a time, but is not safe for simultaneous use by multiple threads.
-    private readonly object _connectionGate = new();
-    private readonly Stack<SqlConnection> _connectionsPool = new();
-
-    private readonly CancellationTokenSource _shutdownTokenSource = new();
-
-    private readonly object _ownershipLock = ownershipLock;
-
-    internal void Initialize(
-        Action<SqlConnection, CancellationToken> initializer,
-        CancellationToken cancellationToken)
+    private readonly struct PooledConnection(SQLitePersistentStorage storage, SqlConnection sqlConnection) : IDisposable
     {
-        // This is our startup path.  No other code can be running.  So it's safe for us to access a connection that
-        // can talk to the db without having to be on the reader/writer scheduler queue.
-        using var _ = GetPooledConnection(checkScheduler: false, out var connection);
+        public readonly SqlConnection Connection = sqlConnection;
 
-        initializer(connection, cancellationToken);
+        public void Dispose()
+            => storage.ReleaseConnection(Connection);
     }
 
     /// <summary>
@@ -45,7 +27,7 @@ internal void Initialize(
     /// longer in use. In particular, make sure to avoid letting a connection lease cross an <see langword="await"/>
     /// boundary, as it will prevent code in the asynchronous operation from using the existing connection.
     /// </remarks>
-    internal PooledConnection GetPooledConnection(out SqlConnection connection)
+    private PooledConnection GetPooledConnection(out SqlConnection connection)
         => GetPooledConnection(checkScheduler: true, out connection);
 
     /// <summary>
@@ -58,8 +40,8 @@ private PooledConnection GetPooledConnection(bool checkScheduler, out SqlConnect
         if (checkScheduler)
         {
             var scheduler = TaskScheduler.Current;
-            if (scheduler != connectionPoolService.Scheduler.ConcurrentScheduler && scheduler != connectionPoolService.Scheduler.ExclusiveScheduler)
-                throw new InvalidOperationException($"Cannot get a connection to the DB unless running on one of {nameof(SQLiteConnectionPoolService)}'s schedulers");
+            if (scheduler != this.Scheduler.ConcurrentScheduler && scheduler != this.Scheduler.ExclusiveScheduler)
+                throw new InvalidOperationException($"Cannot get a connection to the DB unless running on one of {nameof(SQLitePersistentStorage)}'s schedulers");
         }
 
         var result = new PooledConnection(this, GetConnection());
@@ -77,7 +59,7 @@ private SqlConnection GetConnection()
         }
 
         // Otherwise create a new connection.
-        return SqlConnection.Create(faultInjector, databasePath);
+        return SqlConnection.Create(_faultInjector, this.DatabaseFile);
     }
 
     private void ReleaseConnection(SqlConnection connection)

src/Workspaces/Core/Portable/Storage/SQLite/v2/SQLitePersistentStorage.Accessor.cs

@@ -163,13 +163,13 @@ private Optional<T> ReadColumn<T, TData>(
                 // We're reading.  All current scenarios have this happening under the concurrent/read-only scheduler.
                 // If this assert fires either a bug has been introduced, or there is a valid scenario for a writing
                 // codepath to read a column and this assert should be adjusted.
-                Contract.ThrowIfFalse(TaskScheduler.Current == Storage._connectionPoolService.Scheduler.ConcurrentScheduler);
+                Contract.ThrowIfFalse(TaskScheduler.Current == this.Storage.Scheduler.ConcurrentScheduler);
 
                 cancellationToken.ThrowIfCancellationRequested();
 
                 if (!Storage._shutdownTokenSource.IsCancellationRequested)
                 {
-                    using var _ = Storage._connectionPool.GetPooledConnection(out var connection);
+                    using var _ = this.Storage.GetPooledConnection(out var connection);
 
                     // We're in the reading-only scheduler path, so we can't allow TryGetDatabaseId to write.  Note that
                     // this is ok, and actually provides the semantics we want.  Specifically, we can be trying to read
@@ -222,13 +222,13 @@ public Task<bool> WriteStreamAsync(TKey key, string name, Stream stream, Checksu
             private bool WriteStream(TKey key, string dataName, Stream stream, Checksum? checksum, CancellationToken cancellationToken)
             {
                 // We're writing.  This better always be under the exclusive scheduler.
-                Contract.ThrowIfFalse(TaskScheduler.Current == Storage._connectionPoolService.Scheduler.ExclusiveScheduler);
+                Contract.ThrowIfFalse(TaskScheduler.Current == this.Storage.Scheduler.ExclusiveScheduler);
 
                 cancellationToken.ThrowIfCancellationRequested();
 
                 if (!Storage._shutdownTokenSource.IsCancellationRequested)
                 {
-                    using var _ = Storage._connectionPool.GetPooledConnection(out var connection);
+                    using var _ = this.Storage.GetPooledConnection(out var connection);
 
                     // Determine the appropriate data-id to store this stream at.  We already are running
                     // with an exclusive write lock on the DB, so it's safe for us to write the data id to 
@@ -361,7 +361,7 @@ private void InsertOrReplaceBlobIntoWriteCache(
                 ReadOnlySpan<byte> dataBytes)
             {
                 // We're writing.  This better always be under the exclusive scheduler.
-                Contract.ThrowIfFalse(TaskScheduler.Current == Storage._connectionPoolService.Scheduler.ExclusiveScheduler);
+                Contract.ThrowIfFalse(TaskScheduler.Current == this.Storage.Scheduler.ExclusiveScheduler);
 
                 using (var resettableStatement = connection.GetResettableStatement(
                     _insert_or_replace_into_writecache_table_values_0primarykey_1checksum_2data))