Merge pull request #73294 from CyrusNajmabadi/storageLock

src/Tools/IdeBenchmarks/SQLitePersistentStorageBenchmark.cs

@@ -67,7 +67,7 @@ public void GlobalSetup()
             _storageService = new SQLitePersistentStorageService(connectionPoolService, new StorageConfiguration(), asyncListener);
 
             var solution = _workspace.CurrentSolution;
-            _storage = _storageService.GetStorageWorkerAsync(SolutionKey.ToSolutionKey(solution), CancellationToken.None).AsTask().GetAwaiter().GetResult();
+            _storage = _storageService.GetStorageAsync(SolutionKey.ToSolutionKey(solution), CancellationToken.None).AsTask().GetAwaiter().GetResult();
 
             Console.WriteLine("Storage type: " + _storage.GetType());
             _document = _workspace.CurrentSolution.Projects.Single().Documents.Single();

src/VisualStudio/CSharp/Test/PersistentStorage/AbstractPersistentStorageTests.cs

@@ -23,7 +23,7 @@
 namespace Microsoft.CodeAnalysis.UnitTests.WorkspaceServices
 {
     [UseExportProvider]
-    public abstract class AbstractPersistentStorageTests : IDisposable
+    public abstract class AbstractPersistentStorageTests : IAsyncDisposable
     {
         public enum Size
         {
@@ -91,10 +91,10 @@ protected AbstractPersistentStorageTests()
             IPersistentStorageFaultInjector? faultInjector,
             string rootFolder);
 
-        public void Dispose()
+        public async ValueTask DisposeAsync()
         {
             // This should cause the service to release the cached connection it maintains for the primary workspace
-            _storageService?.GetTestAccessor().Shutdown();
+            await (_storageService?.GetTestAccessor().ShutdownAsync() ?? Task.CompletedTask);
             _persistentFolderRoot.Dispose();
         }
 
@@ -1013,7 +1013,7 @@ protected Solution CreateOrOpenSolution(TempDirectory? persistentFolder = null,
         {
             // If we handed out one for a previous test, we need to shut that down first
             persistentFolder ??= _persistentFolder;
-            _storageService?.GetTestAccessor().Shutdown();
+            await (_storageService?.GetTestAccessor().ShutdownAsync() ?? Task.CompletedTask);
             var configuration = new MockPersistentStorageConfiguration(solution.Id, persistentFolder.Path, throwOnFailure);
 
             _storageService = GetStorageService(solution.Workspace.Services.SolutionServices.ExportProvider, configuration, faultInjector, _persistentFolder.Path);
@@ -1022,7 +1022,7 @@ protected Solution CreateOrOpenSolution(TempDirectory? persistentFolder = null,
             // If we're injecting faults, we expect things to be strange
             if (faultInjector == null)
             {
-                Assert.NotEqual(NoOpPersistentStorage.TestAccessor.StorageInstance, storage);
+                Assert.NotEqual(NoOpPersistentStorage.TestAccessor.GetStorageInstance(SolutionKey.ToSolutionKey(solution)), storage);
             }
 
             return storage;
@@ -1032,7 +1032,7 @@ protected Solution CreateOrOpenSolution(TempDirectory? persistentFolder = null,
             HostWorkspaceServices services, SolutionKey solutionKey, IPersistentStorageFaultInjector? faultInjector = null)
         {
             // If we handed out one for a previous test, we need to shut that down first
-            _storageService?.GetTestAccessor().Shutdown();
+            await (_storageService?.GetTestAccessor().ShutdownAsync() ?? Task.CompletedTask);
             var configuration = new MockPersistentStorageConfiguration(solutionKey.Id, _persistentFolder.Path, throwOnFailure: true);
 
             _storageService = GetStorageService(services.SolutionServices.ExportProvider, configuration, faultInjector, _persistentFolder.Path);
@@ -1041,7 +1041,7 @@ protected Solution CreateOrOpenSolution(TempDirectory? persistentFolder = null,
             // If we're injecting faults, we expect things to be strange
             if (faultInjector == null)
             {
-                Assert.NotEqual(NoOpPersistentStorage.TestAccessor.StorageInstance, storage);
+                Assert.NotEqual(NoOpPersistentStorage.TestAccessor.GetStorageInstance(solutionKey), storage);
             }
 
             return storage;

src/VisualStudio/CSharp/Test/PersistentStorage/SQLiteV2PersistentStorageTests.cs

@@ -8,7 +8,6 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
-using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.SQLite.v2;
 using Microsoft.CodeAnalysis.Storage;

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

@@ -17,19 +17,15 @@ namespace Microsoft.CodeAnalysis.Storage;
 /// A service that enables storing and retrieving of information associated with solutions,
 /// projects or documents across runtime sessions.
 /// </summary>
-internal abstract partial class AbstractPersistentStorageService : IChecksummedPersistentStorageService
+internal abstract partial class AbstractPersistentStorageService(IPersistentStorageConfiguration configuration) : IChecksummedPersistentStorageService
 {
-    protected readonly IPersistentStorageConfiguration Configuration;
+    protected readonly IPersistentStorageConfiguration Configuration = configuration;
 
     /// <summary>
     /// This lock guards all mutable fields in this type.
     /// </summary>
     private readonly SemaphoreSlim _lock = new(initialCount: 1);
     private ReferenceCountedDisposable<IChecksummedPersistentStorage>? _currentPersistentStorage;
-    private SolutionId? _currentPersistentStorageSolutionId;
-
-    protected AbstractPersistentStorageService(IPersistentStorageConfiguration configuration)
-        => Configuration = configuration;
 
     protected abstract string GetDatabaseFilePath(string workingFolderPath);
 
@@ -41,57 +37,58 @@ protected AbstractPersistentStorageService(IPersistentStorageConfiguration confi
     protected abstract ValueTask<IChecksummedPersistentStorage?> TryOpenDatabaseAsync(SolutionKey solutionKey, string workingFolderPath, string databaseFilePath, CancellationToken cancellationToken);
     protected abstract bool ShouldDeleteDatabase(Exception exception);
 
-    public ValueTask<IChecksummedPersistentStorage> GetStorageAsync(SolutionKey solutionKey, CancellationToken cancellationToken)
+    public async ValueTask<IChecksummedPersistentStorage> GetStorageAsync(SolutionKey solutionKey, CancellationToken cancellationToken)
     {
-        return solutionKey.FilePath == null
-            ? new(NoOpPersistentStorage.GetOrThrow(Configuration.ThrowOnFailure))
-            : GetStorageWorkerAsync(solutionKey, cancellationToken);
-    }
+        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.
+        using var existing = _currentPersistentStorage?.TryAddReference();
+        if (solutionKey == existing?.Target.SolutionKey)
+        {
+            // Success, we can use the current storage system.  Ensure we increment the reference count again, so that
+            // this stays alive for the caller when the above reference count drops.
+            return PersistentStorageReferenceCountedDisposableWrapper.AddReferenceCountToAndCreateWrapper(existing);
+        }
+
+        var workingFolder = Configuration.TryGetStorageLocation(solutionKey);
+        if (workingFolder == null)
+            return NoOpPersistentStorage.GetOrThrow(solutionKey, Configuration.ThrowOnFailure);
 
-    internal async ValueTask<IChecksummedPersistentStorage> GetStorageWorkerAsync(SolutionKey solutionKey, CancellationToken cancellationToken)
-    {
         using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
         {
-            // Do we already have storage for this?
-            if (solutionKey.Id == _currentPersistentStorageSolutionId)
+            // See if another thread set to the solution we care about while we were waiting on the lock.
+            using var current = _currentPersistentStorage?.TryAddReference();
+            if (solutionKey != current?.Target.SolutionKey)
             {
-                // We do, great. Increment our ref count for our caller.  They'll decrement it
-                // when done with it.
-                return PersistentStorageReferenceCountedDisposableWrapper.AddReferenceCountToAndCreateWrapper(_currentPersistentStorage!);
+                // We either don't have a current storage, or the current storage was pointing at some other solution. In
+                // either case, we want to create a new storage instance and point at that.  If the current storage was
+                // point at some other solution, then lower its ref count (by calling DisposeAsync, outside of the lock) to
+                // indicate that *we* are letting go of it.
+                _ = DisposeStorageAsync(_currentPersistentStorage);
+
+                // Create and cache a new storage instance associated with this particular solution.
+                // It will initially have a ref-count of 1 due to our reference to it.
+                _currentPersistentStorage = new ReferenceCountedDisposable<IChecksummedPersistentStorage>(
+                    await CreatePersistentStorageAsync(solutionKey, workingFolder, cancellationToken).ConfigureAwait(false));
             }
 
-            var workingFolder = Configuration.TryGetStorageLocation(solutionKey);
-            if (workingFolder == null)
-                return NoOpPersistentStorage.GetOrThrow(Configuration.ThrowOnFailure);
+            // Now increment the reference count and return to our caller.  The current ref count for this instance will
+            // be at least 2.  Until all the callers *and* us decrement the refcounts, this instance will not be
+            // actually disposed.
+            Contract.ThrowIfNull(_currentPersistentStorage);
+            return PersistentStorageReferenceCountedDisposableWrapper.AddReferenceCountToAndCreateWrapper(_currentPersistentStorage);
+        }
 
-            // If we already had some previous cached service, let's let it start cleaning up
-            if (_currentPersistentStorage != null)
+        static async Task DisposeStorageAsync(ReferenceCountedDisposable<IChecksummedPersistentStorage>? storage)
+        {
+            if (storage != null)
             {
-                var storageToDispose = _currentPersistentStorage;
-
-                // Kick off a task to actually go dispose the previous cached storage instance.
-                // This will remove the single ref count we ourselves added when we cached the
-                // instance.  Then once all other existing clients who are holding onto this
-                // instance let go, it will finally get truly disposed.
-                // This operation is not safe to cancel (as dispose must happen).
-                _ = Task.Run(storageToDispose.Dispose, CancellationToken.None);
-
-                _currentPersistentStorage = null;
-                _currentPersistentStorageSolutionId = null;
+                // Intentionally yield, so we don't execute any of this code outside the outer lock.
+                await Task.Yield().ConfigureAwait(false);
+                await storage.DisposeAsync().ConfigureAwait(false);
             }
-
-            var storage = await CreatePersistentStorageAsync(solutionKey, workingFolder, cancellationToken).ConfigureAwait(false);
-            Contract.ThrowIfNull(storage);
-
-            // Create and cache a new storage instance associated with this particular solution.
-            // It will initially have a ref-count of 1 due to our reference to it.
-            _currentPersistentStorage = new ReferenceCountedDisposable<IChecksummedPersistentStorage>(storage);
-            _currentPersistentStorageSolutionId = solutionKey.Id;
-
-            // Now increment the reference count and return to our caller.  The current ref
-            // count for this instance will be 2.  Until all the callers *and* us decrement
-            // the refcounts, this instance will not be actually disposed.
-            return PersistentStorageReferenceCountedDisposableWrapper.AddReferenceCountToAndCreateWrapper(_currentPersistentStorage);
         }
     }
 
@@ -108,7 +105,7 @@ internal async ValueTask<IChecksummedPersistentStorage> GetStorageWorkerAsync(So
         if (result != null)
             return result;
 
-        return NoOpPersistentStorage.GetOrThrow(Configuration.ThrowOnFailure);
+        return NoOpPersistentStorage.GetOrThrow(solutionKey, Configuration.ThrowOnFailure);
     }
 
     private async ValueTask<IChecksummedPersistentStorage?> TryCreatePersistentStorageAsync(
@@ -147,30 +144,30 @@ bool Recover(Exception ex)
         }
     }
 
-    private void Shutdown()
+    private async Task ShutdownAsync(CancellationToken cancellationToken)
     {
         ReferenceCountedDisposable<IChecksummedPersistentStorage>? storage = null;
 
-        lock (_lock)
+        using (await _lock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
         {
             // We will transfer ownership in a thread-safe way out so we can dispose outside the lock
             storage = _currentPersistentStorage;
             _currentPersistentStorage = null;
-            _currentPersistentStorageSolutionId = null;
         }
 
         // Dispose storage outside of the lock. Note this only removes our reference count; clients who are still
         // using this will still be holding a reference count.
-        storage?.Dispose();
+        if (storage != null)
+            await storage.DisposeAsync().ConfigureAwait(false);
     }
 
     internal TestAccessor GetTestAccessor()
         => new(this);
 
     internal readonly struct TestAccessor(AbstractPersistentStorageService service)
     {
-        public void Shutdown()
-            => service.Shutdown();
+        public Task ShutdownAsync()
+            => service.ShutdownAsync(CancellationToken.None);
     }
 
     /// <summary>
@@ -197,6 +194,9 @@ public void Dispose()
         public ValueTask DisposeAsync()
             => _storage.DisposeAsync();
 
+        public SolutionKey SolutionKey
+            => _storage.Target.SolutionKey;
+
         public Task<bool> ChecksumMatchesAsync(string name, Checksum checksum, CancellationToken cancellationToken)
             => _storage.Target.ChecksumMatchesAsync(name, checksum, cancellationToken);
 

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

@@ -36,5 +36,5 @@ public LegacyPersistentStorageService()
     }
 
     public IPersistentStorage GetStorage(Solution solution)
-        => NoOpPersistentStorage.GetOrThrow(throwOnFailure: false);
+        => NoOpPersistentStorage.GetOrThrow(SolutionKey.ToSolutionKey(solution), throwOnFailure: false);
 }

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

@@ -25,7 +25,6 @@ internal sealed partial class SQLitePersistentStorage : AbstractPersistentStorag
 {
     private readonly CancellationTokenSource _shutdownTokenSource = new();
 
-    private readonly SolutionKey _solutionKey;
     private readonly string _solutionDirectory;
 
     private readonly SQLiteConnectionPoolService _connectionPoolService;
@@ -54,10 +53,9 @@ internal sealed partial class SQLitePersistentStorage : AbstractPersistentStorag
         string databaseFile,
         IAsynchronousOperationListener asyncListener,
         IPersistentStorageFaultInjector? faultInjector)
-        : base(workingFolderPath, solutionKey.FilePath!, databaseFile)
+        : base(solutionKey, workingFolderPath, databaseFile)
     {
         Contract.ThrowIfNull(solutionKey.FilePath);
-        _solutionKey = solutionKey;
         _solutionDirectory = PathUtilities.GetDirectoryName(solutionKey.FilePath);
         _connectionPoolService = connectionPoolService;
 

src/Workspaces/Core/Portable/Storage/SQLite/v2/SQLitePersistentStorageService.cs

@@ -89,7 +89,7 @@ protected override string GetDatabaseFilePath(string workingFolderPath)
         }
 
         if (solutionKey.FilePath == null)
-            return new(NoOpPersistentStorage.GetOrThrow(Configuration.ThrowOnFailure));
+            return new(NoOpPersistentStorage.GetOrThrow(solutionKey, Configuration.ThrowOnFailure));
 
         return new(SQLitePersistentStorage.TryCreate(
             connectionPoolService,

src/Workspaces/Core/Portable/Storage/SQLite/v2/SQLitePersistentStorage_SolutionSerialization.cs

@@ -16,13 +16,13 @@ namespace Microsoft.CodeAnalysis.SQLite.v2;
 internal partial class SQLitePersistentStorage
 {
     public override Task<bool> ChecksumMatchesAsync(string name, Checksum checksum, CancellationToken cancellationToken)
-        => _solutionAccessor.ChecksumMatchesAsync(_solutionKey, name, checksum, cancellationToken);
+        => _solutionAccessor.ChecksumMatchesAsync(this.SolutionKey, name, checksum, cancellationToken);
 
     public override Task<Stream?> ReadStreamAsync(string name, Checksum? checksum, CancellationToken cancellationToken)
-        => _solutionAccessor.ReadStreamAsync(_solutionKey, name, checksum, cancellationToken);
+        => _solutionAccessor.ReadStreamAsync(this.SolutionKey, name, checksum, cancellationToken);
 
     public override Task<bool> WriteStreamAsync(string name, Stream stream, Checksum? checksum, CancellationToken cancellationToken)
-        => _solutionAccessor.WriteStreamAsync(_solutionKey, name, stream, checksum, cancellationToken);
+        => _solutionAccessor.WriteStreamAsync(this.SolutionKey, name, stream, checksum, cancellationToken);
 
     private readonly record struct SolutionPrimaryKey();
 

src/Workspaces/Core/Portable/Workspace/Host/PersistentStorage/AbstractPersistentStorage.cs

@@ -13,21 +13,22 @@ namespace Microsoft.CodeAnalysis.Host;
 
 internal abstract class AbstractPersistentStorage : IChecksummedPersistentStorage
 {
+    public SolutionKey SolutionKey { get; }
+
     public string WorkingFolderPath { get; }
-    public string SolutionFilePath { get; }
 
     public string DatabaseFile { get; }
     public string DatabaseDirectory => Path.GetDirectoryName(DatabaseFile) ?? throw ExceptionUtilities.UnexpectedValue(DatabaseFile);
 
     private bool _isDisabled;
 
     protected AbstractPersistentStorage(
+        SolutionKey solutionKey,
         string workingFolderPath,
-        string solutionFilePath,
         string databaseFile)
     {
+        this.SolutionKey = solutionKey;
         this.WorkingFolderPath = workingFolderPath;
-        this.SolutionFilePath = solutionFilePath;
         this.DatabaseFile = databaseFile;
 
         if (!Directory.Exists(this.DatabaseDirectory))

src/Workspaces/Core/Portable/Workspace/Host/PersistentStorage/IChecksummedPersistentStorage.cs

@@ -12,6 +12,11 @@ namespace Microsoft.CodeAnalysis.Host;
 
 internal interface IChecksummedPersistentStorage : IPersistentStorage
 {
+    /// <summary>
+    /// The solution this is a storage instance for.
+    /// </summary>
+    SolutionKey SolutionKey { get; }
+
     /// <summary>
     /// <see langword="true"/> if the data we have for the solution with the given <paramref name="name"/> has the
     /// provided <paramref name="checksum"/>.

src/Workspaces/Core/Portable/Workspace/Host/PersistentStorage/NoOpPersistentStorage.cs

@@ -11,27 +11,21 @@
 
 namespace Microsoft.CodeAnalysis.Host;
 
-internal class NoOpPersistentStorage : IChecksummedPersistentStorage
+internal class NoOpPersistentStorage(SolutionKey solutionKey) : IChecksummedPersistentStorage
 {
-    private static readonly IChecksummedPersistentStorage Instance = new NoOpPersistentStorage();
+    public SolutionKey SolutionKey => solutionKey;
 
-    private NoOpPersistentStorage()
-    {
-    }
-
-    public static IChecksummedPersistentStorage GetOrThrow(bool throwOnFailure)
+    public static IChecksummedPersistentStorage GetOrThrow(SolutionKey solutionKey, bool throwOnFailure)
         => throwOnFailure
             ? throw new InvalidOperationException("Database was not supported")
-            : Instance;
+            : new NoOpPersistentStorage(solutionKey);
 
     public void Dispose()
     {
     }
 
     public ValueTask DisposeAsync()
-    {
-        return ValueTaskFactory.CompletedTask;
-    }
+        => ValueTaskFactory.CompletedTask;
 
     public Task<bool> ChecksumMatchesAsync(string name, Checksum checksum, CancellationToken cancellationToken)
         => SpecializedTasks.False;
@@ -98,6 +92,6 @@ public Task<bool> WriteStreamAsync(DocumentKey documentKey, string name, Stream
 
     public readonly struct TestAccessor
     {
-        public static readonly IChecksummedPersistentStorage StorageInstance = Instance;
+        public static IChecksummedPersistentStorage GetStorageInstance(SolutionKey solutionKey) => new NoOpPersistentStorage(solutionKey);
     }
 }