Synchronise language service project updates

src/Microsoft.VisualStudio.ProjectSystem.Managed.VS/ProjectSystem/VS/LanguageServices/Workspace.cs

@@ -156,27 +156,24 @@ public void ChainDisposal(IDisposable disposable)
     /// both evaluation and build updates may arrive in any order, so long as values
     /// of each type are ordered correctly.
     /// </para>
-    /// <para>
-    /// Calls must not overlap. This method is not thread-safe. This method is designed
-    /// to be called from a dataflow ActionBlock, which will serialize calls, so we
-    /// needn't perform any locking or protection here.
-    /// </para>
     /// </remarks>
     /// <param name="update">The project update to integrate.</param>
     /// <returns>A task that completes when the update has been integrated.</returns>
-    internal async Task OnWorkspaceUpdateAsync(IProjectVersionedValue<WorkspaceUpdate> update)
+    internal Task OnWorkspaceUpdateAsync(IProjectVersionedValue<WorkspaceUpdate> update)
+    {
+        // Prevent disposal during the update.
+        return ExecuteUnderLockAsync(async token =>
         {
-        Verify.NotDisposed(this);
-
             await InitializeAsync(_unloadCancellationToken);
 
             Assumes.True(_state is WorkspaceState.Uninitialized or WorkspaceState.Initialized);
 
-        await _joinableTaskFactory.RunAsync(
-            async () =>
+            await _joinableTaskFactory.RunAsync(ApplyUpdateWithinLockAsync);
+        });
+
+        async Task ApplyUpdateWithinLockAsync()
         {
-                // Calls never overlap. No synchronisation is needed here.
-                // We can receive either evaluation OR build data first.
+            // We will always receive an evaluation update before the first build update.
 
             if (TryTransition(WorkspaceState.Uninitialized, WorkspaceState.Initialized))
             {
@@ -215,7 +212,7 @@ await _joinableTaskFactory.RunAsync(
                 // to produce an upstream NFE.
                 throw;
             }
-            });
+        }
     }
 
     private async Task OnEvaluationUpdateAsync(IProjectVersionedValue<EvaluationUpdate> evaluationUpdate)
@@ -541,28 +538,34 @@ private static IComparable GetConfiguredProjectVersion(IProjectValueVersions upd
         return update.DataSourceVersions[ProjectDataSources.ConfiguredProjectVersion];
     }
 
-    public async Task WriteAsync(Func<IWorkspace, Task> action, CancellationToken cancellationToken)
+    public Task WriteAsync(Func<IWorkspace, Task> action, CancellationToken cancellationToken)
     {
         Requires.NotNull(action);
-        Verify.NotDisposed(this);
 
         cancellationToken = CancellationTokenExtensions.CombineWith(_unloadCancellationToken, cancellationToken).Token;
 
+        return ExecuteUnderLockAsync(async _ =>
+            {
                 await WhenContextCreated(cancellationToken);
 
-        await ExecuteUnderLockAsync(_ => action(this), cancellationToken);
+                await action(this);
+            },
+            cancellationToken);
     }
 
-    public async Task<T> WriteAsync<T>(Func<IWorkspace, Task<T>> action, CancellationToken cancellationToken)
+    public Task<T> WriteAsync<T>(Func<IWorkspace, Task<T>> action, CancellationToken cancellationToken)
     {
         Requires.NotNull(action);
-        Verify.NotDisposed(this);
 
         cancellationToken = CancellationTokenExtensions.CombineWith(_unloadCancellationToken, cancellationToken).Token;
 
+        return ExecuteUnderLockAsync(async _ =>
+            {
                 await WhenContextCreated(cancellationToken);
 
-        return await ExecuteUnderLockAsync(_ => action(this), cancellationToken);
+                return await action(this);
+            },
+            cancellationToken);
     }
 
     private async Task WhenContextCreated(CancellationToken cancellationToken)
@@ -571,7 +574,7 @@ private async Task WhenContextCreated(CancellationToken cancellationToken)
 
         // Join the same collection that's used by our dataflow nodes, so that if we are called on
         // the main thread, we don't block anything that might prohibit dataflow from progressing
-        // this workspace's initialisation (leading to deadlock).
+        // this workspace's initialization (leading to deadlock).
         using (_joinableTaskCollection.Join())
         {
             // Ensure we have received enough data to create the context.

tests/Microsoft.VisualStudio.ProjectSystem.Managed.VS.UnitTests/ProjectSystem/VS/LanguageServices/WorkspaceTests.cs

@@ -85,11 +85,20 @@ public async Task Dispose_TriggersObjectDisposedExceptionsOnPublicMembers()
 
         await Assert.ThrowsAsync<ObjectDisposedException>(() => workspace.WriteAsync(w => Task.CompletedTask, CancellationToken.None));
         await Assert.ThrowsAsync<ObjectDisposedException>(() => workspace.WriteAsync(w => TaskResult.EmptyString, CancellationToken.None));
-        await Assert.ThrowsAsync<ObjectDisposedException>(() => workspace.OnWorkspaceUpdateAsync(null!));
 
         Assert.Throws<ObjectDisposedException>(() => workspace.ChainDisposal(null!));
     }
 
+    [Fact]
+    public async Task Dispose_TriggersOperationCanceledExceptionOnWorkspaceUpdates()
+    {
+        var workspace = await CreateInstanceAsync();
+
+        await workspace.DisposeAsync();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(() => workspace.OnWorkspaceUpdateAsync(null!));
+    }
+
     [Theory]
     [CombinatorialData]
     public async Task WriteAsync_ThrowsIfNullAction(bool isGeneric)