From 18feba09828bc498e2659c30c6bff48553a0281a Mon Sep 17 00:00:00 2001
From: Jason Malinowski <jason.malinowski@microsoft.com>
Date: Fri, 24 Feb 2023 16:11:25 -0800
Subject: [PATCH] Fix marshalability of IWorkspaceProject.StartBatchAsync

We have to make this return an interface marked with [RpcMarshalable],
and since that can't really use IAsyncDisposable to say when it's done,
I'm adding an explicit Apply method to make up for it.
---
 .../BrokeredService/WorkspaceProject.cs       | 27 +++++++++++++++++--
 ...soft.CodeAnalysis.Remote.Workspaces.csproj |  1 +
 .../Core/ProjectSystem/IWorkspaceProject.cs   |  2 +-
 .../ProjectSystem/IWorkspaceProjectBatch.cs   | 16 +++++++++++
 4 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProjectBatch.cs

diff --git a/src/VisualStudio/Core/Def/ProjectSystem/BrokeredService/WorkspaceProject.cs b/src/VisualStudio/Core/Def/ProjectSystem/BrokeredService/WorkspaceProject.cs
index 19fc70556a65c..f789f0eacd8b6 100644
--- a/src/VisualStudio/Core/Def/ProjectSystem/BrokeredService/WorkspaceProject.cs
+++ b/src/VisualStudio/Core/Def/ProjectSystem/BrokeredService/WorkspaceProject.cs
@@ -159,9 +159,32 @@ public Task SetProjectHasAllInformationAsync(bool hasAllInformation, Cancellatio
             return Task.CompletedTask;
         }
 
-        public Task<IAsyncDisposable> StartBatchAsync(CancellationToken cancellationToken)
+        public Task<IWorkspaceProjectBatch> StartBatchAsync(CancellationToken cancellationToken)
         {
-            return Task.FromResult(_project.CreateBatchScope());
+            return Task.FromResult<IWorkspaceProjectBatch>(new WorkspaceProjectBatch(_project.CreateBatchScope()));
+        }
+
+        private class WorkspaceProjectBatch : IWorkspaceProjectBatch
+        {
+            private IAsyncDisposable? _batch;
+
+            public WorkspaceProjectBatch(IAsyncDisposable batch)
+            {
+                _batch = batch;
+            }
+
+            public async Task ApplyAsync(CancellationToken cancellationToken)
+            {
+                if (_batch == null)
+                    throw new InvalidOperationException("The batch has already been applied.");
+
+                await _batch.DisposeAsync().ConfigureAwait(false);
+                _batch = null;
+            }
+
+            public void Dispose()
+            {
+            }
         }
     }
 }
diff --git a/src/Workspaces/Remote/Core/Microsoft.CodeAnalysis.Remote.Workspaces.csproj b/src/Workspaces/Remote/Core/Microsoft.CodeAnalysis.Remote.Workspaces.csproj
index 23ee18bb69462..a792c6aee9825 100644
--- a/src/Workspaces/Remote/Core/Microsoft.CodeAnalysis.Remote.Workspaces.csproj
+++ b/src/Workspaces/Remote/Core/Microsoft.CodeAnalysis.Remote.Workspaces.csproj
@@ -46,6 +46,7 @@
     <RestrictedInternalsVisibleTo Include="Microsoft.CodeAnalysis.UnitTesting.SourceBasedTestDiscovery" Partner="UnitTesting" Key="$(UnitTestingKey)" />
     <RestrictedInternalsVisibleTo Include="Microsoft.CodeAnalysis.UnitTesting.SourceBasedTestDiscovery.Core" Partner="UnitTesting" Key="$(UnitTestingKey)" />
     <InternalsVisibleTo Include="Microsoft.CodeAnalysis.LanguageServer" Key="$(MicrosoftCodeAnalysisLanguageServerKey)" />
+    <InternalsVisibleTo Include="Microsoft.CodeAnalysis.LanguageServer.UnitTests" Key="$(MicrosoftCodeAnalysisLanguageServerKey)" />
     <InternalsVisibleTo Include="Roslyn.VisualStudio.DiagnosticsWindow" />
     <InternalsVisibleTo Include="Microsoft.VisualStudio.LanguageServices" />
     <InternalsVisibleTo Include="Microsoft.VisualStudio.LanguageServices.CodeLens" />
diff --git a/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProject.cs b/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProject.cs
index c96554fdd830f..4e6ab5f9fbefc 100644
--- a/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProject.cs
+++ b/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProject.cs
@@ -38,5 +38,5 @@ internal interface IWorkspaceProject : IDisposable
 
     Task SetProjectHasAllInformationAsync(bool hasAllInformation, CancellationToken cancellationToken);
 
-    Task<IAsyncDisposable> StartBatchAsync(CancellationToken cancellationToken);
+    Task<IWorkspaceProjectBatch> StartBatchAsync(CancellationToken cancellationToken);
 }
diff --git a/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProjectBatch.cs b/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProjectBatch.cs
new file mode 100644
index 0000000000000..e8eac3ee22ea0
--- /dev/null
+++ b/src/Workspaces/Remote/Core/ProjectSystem/IWorkspaceProjectBatch.cs
@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using StreamJsonRpc;
+
+namespace Microsoft.CodeAnalysis.Remote.ProjectSystem;
+
+[RpcMarshalable]
+internal interface IWorkspaceProjectBatch : IDisposable
+{
+    Task ApplyAsync(CancellationToken cancellationToken);
+}