Keep storage objects of serialized metadata references alive during solution replication (#49974)

src/VisualStudio/Core/Test.Next/Mocks/SimpleAssetSource.cs

@@ -36,9 +36,11 @@ public SimpleAssetSource(ISerializerService serializerService, IReadOnlyDictiona
                 if (_map.TryGetValue(checksum, out var data))
                 {
                     using var stream = new MemoryStream();
+                    using var context = SolutionReplicationContext.Create();
+
                     using (var writer = new ObjectWriter(stream, leaveOpen: true, cancellationToken))
                     {
-                        _serializerService.Serialize(data, writer, cancellationToken);
+                        _serializerService.Serialize(data, writer, context, cancellationToken);
                     }
 
                     stream.Position = 0;

src/VisualStudio/Core/Test.Next/Remote/SerializationValidator.cs

@@ -82,10 +82,11 @@ public async Task<T> GetValueAsync<T>(Checksum checksum)
             var data = (await AssetStorage.GetTestAccessor().GetAssetAsync(checksum, CancellationToken.None).ConfigureAwait(false))!;
             Contract.ThrowIfNull(data.Value);
 
+            using var context = SolutionReplicationContext.Create();
             using var stream = SerializableBytes.CreateWritableStream();
             using (var writer = new ObjectWriter(stream, leaveOpen: true))
             {
-                Serializer.Serialize(data.Value, writer, CancellationToken.None);
+                Serializer.Serialize(data.Value, writer, context, CancellationToken.None);
             }
 
             stream.Position = 0;

src/VisualStudio/Core/Test.Next/Remote/SnapshotSerializationTests.cs

@@ -655,9 +655,11 @@ public void TestEncodingSerialization()
             var sourceText = SourceText.From("Hello", Encoding.UTF8);
             using (var stream = SerializableBytes.CreateWritableStream())
             {
+                using var context = SolutionReplicationContext.Create();
+
                 using (var objectWriter = new ObjectWriter(stream, leaveOpen: true))
                 {
-                    serializer.Serialize(sourceText, objectWriter, CancellationToken.None);
+                    serializer.Serialize(sourceText, objectWriter, context, CancellationToken.None);
                 }
 
                 stream.Position = 0;
@@ -672,9 +674,11 @@ public void TestEncodingSerialization()
             sourceText = SourceText.From("Hello", new NotSerializableEncoding());
             using (var stream = SerializableBytes.CreateWritableStream())
             {
+                using var context = SolutionReplicationContext.Create();
+
                 using (var objectWriter = new ObjectWriter(stream, leaveOpen: true))
                 {
-                    serializer.Serialize(sourceText, objectWriter, CancellationToken.None);
+                    serializer.Serialize(sourceText, objectWriter, context, CancellationToken.None);
                 }
 
                 stream.Position = 0;
@@ -701,9 +705,11 @@ public void TestCompilationOptions_NullableAndImport()
             void VerifyOptions(CompilationOptions originalOptions)
             {
                 using var stream = SerializableBytes.CreateWritableStream();
+                using var context = SolutionReplicationContext.Create();
+
                 using (var objectWriter = new ObjectWriter(stream, leaveOpen: true))
                 {
-                    serializer.Serialize(originalOptions, objectWriter, CancellationToken.None);
+                    serializer.Serialize(originalOptions, objectWriter, context, CancellationToken.None);
                 }
 
                 stream.Position = 0;
@@ -750,10 +756,11 @@ private static async Task VerifyOptionSetsAsync(Workspace workspace, Action<Opti
         private static SolutionAsset CloneAsset(ISerializerService serializer, SolutionAsset asset)
         {
             using var stream = SerializableBytes.CreateWritableStream();
+            using var context = SolutionReplicationContext.Create();
 
             using (var writer = new ObjectWriter(stream, leaveOpen: true))
             {
-                serializer.Serialize(asset.Value, writer, CancellationToken.None);
+                serializer.Serialize(asset.Value, writer, context, CancellationToken.None);
             }
 
             stream.Position = 0;

src/Workspaces/Core/Portable/Remote/ISerializerService.cs

@@ -12,9 +12,9 @@ namespace Microsoft.CodeAnalysis.Serialization
 {
     internal interface ISerializerService : IWorkspaceService
     {
-        void Serialize(object value, ObjectWriter writer, CancellationToken cancellationToken);
+        void Serialize(object value, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken);
 
-        void SerializeSourceText(SerializableSourceText text, ObjectWriter writer, CancellationToken cancellationToken);
+        void SerializeSourceText(SerializableSourceText text, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken);
 
         void SerializeCompilationOptions(CompilationOptions options, ObjectWriter writer, CancellationToken cancellationToken);
 
@@ -22,7 +22,7 @@ internal interface ISerializerService : IWorkspaceService
 
         void SerializeProjectReference(ProjectReference reference, ObjectWriter writer, CancellationToken cancellationToken);
 
-        void SerializeMetadataReference(MetadataReference reference, ObjectWriter writer, CancellationToken cancellationToken);
+        void SerializeMetadataReference(MetadataReference reference, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken);
 
         void SerializeAnalyzerReference(AnalyzerReference reference, ObjectWriter writer, CancellationToken cancellationToken);
 

src/Workspaces/Core/Portable/Serialization/SerializerService.cs

@@ -101,7 +101,7 @@ public Checksum CreateChecksum(object value, CancellationToken cancellationToken
             }
         }
 
-        public void Serialize(object value, ObjectWriter writer, CancellationToken cancellationToken)
+        public void Serialize(object value, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             var kind = value.GetWellKnownSynchronizationKind();
 
@@ -140,7 +140,7 @@ public void Serialize(object value, ObjectWriter writer, CancellationToken cance
                         return;
 
                     case WellKnownSynchronizationKind.MetadataReference:
-                        SerializeMetadataReference((MetadataReference)value, writer, cancellationToken);
+                        SerializeMetadataReference((MetadataReference)value, writer, context, cancellationToken);
                         return;
 
                     case WellKnownSynchronizationKind.AnalyzerReference:
@@ -148,11 +148,11 @@ public void Serialize(object value, ObjectWriter writer, CancellationToken cance
                         return;
 
                     case WellKnownSynchronizationKind.SerializableSourceText:
-                        SerializeSourceText((SerializableSourceText)value, writer, cancellationToken);
+                        SerializeSourceText((SerializableSourceText)value, writer, context, cancellationToken);
                         return;
 
                     case WellKnownSynchronizationKind.SourceText:
-                        SerializeSourceText(new SerializableSourceText((SourceText)value), writer, cancellationToken);
+                        SerializeSourceText(new SerializableSourceText((SourceText)value), writer, context, cancellationToken);
                         return;
 
                     case WellKnownSynchronizationKind.OptionSet:

src/Workspaces/Core/Portable/Serialization/SerializerService_Asset.cs

@@ -19,11 +19,13 @@ namespace Microsoft.CodeAnalysis.Serialization
     /// </summary>
     internal partial class SerializerService
     {
-        public void SerializeSourceText(SerializableSourceText text, ObjectWriter writer, CancellationToken cancellationToken)
+        public void SerializeSourceText(SerializableSourceText text, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
             if (text.Storage is not null)
             {
+                context.AddResource(text.Storage);
+
                 writer.WriteInt32((int)text.Storage.ChecksumAlgorithm);
                 writer.WriteEncoding(text.Storage.Encoding);
 
@@ -146,10 +148,10 @@ private static ProjectReference DeserializeProjectReference(ObjectReader reader,
             return new ProjectReference(projectId, aliases.ToImmutableArrayOrEmpty(), embedInteropTypes);
         }
 
-        public void SerializeMetadataReference(MetadataReference reference, ObjectWriter writer, CancellationToken cancellationToken)
+        public void SerializeMetadataReference(MetadataReference reference, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
-            WriteMetadataReferenceTo(reference, writer, cancellationToken);
+            WriteMetadataReferenceTo(reference, writer, context, cancellationToken);
         }
 
         private MetadataReference DeserializeMetadataReference(ObjectReader reader, CancellationToken cancellationToken)

src/Workspaces/Core/Portable/Serialization/SerializerService_Reference.cs

@@ -61,13 +61,13 @@ public static Checksum CreateChecksum(AnalyzerReference reference, CancellationT
             return Checksum.Create(stream);
         }
 
-        public virtual void WriteMetadataReferenceTo(MetadataReference reference, ObjectWriter writer, CancellationToken cancellationToken)
+        public virtual void WriteMetadataReferenceTo(MetadataReference reference, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             if (reference is PortableExecutableReference portable)
             {
                 if (portable is ISupportTemporaryStorage supportTemporaryStorage)
                 {
-                    if (TryWritePortableExecutableReferenceBackedByTemporaryStorageTo(supportTemporaryStorage, writer, cancellationToken))
+                    if (TryWritePortableExecutableReferenceBackedByTemporaryStorageTo(supportTemporaryStorage, writer, context, cancellationToken))
                     {
                         return;
                     }
@@ -315,7 +315,7 @@ private static void WriteTo(Metadata? metadata, ObjectWriter writer, Cancellatio
         }
 
         private static bool TryWritePortableExecutableReferenceBackedByTemporaryStorageTo(
-            ISupportTemporaryStorage reference, ObjectWriter writer, CancellationToken cancellationToken)
+            ISupportTemporaryStorage reference, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             var storages = reference.GetStorages();
             if (storages == null)
@@ -328,11 +328,13 @@ private static void WriteTo(Metadata? metadata, ObjectWriter writer, Cancellatio
 
             foreach (var storage in storages)
             {
-                if (!(storage is ITemporaryStorageWithName storage2))
+                if (storage is not ITemporaryStorageWithName storage2)
                 {
                     return false;
                 }
 
+                context.AddResource(storage);
+
                 pooled.Object.Add((storage2.Name, storage2.Offset, storage2.Size));
             }
 

src/Workspaces/Core/Portable/Serialization/SolutionReplicationContext.cs

+// 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 Microsoft.CodeAnalysis.PooledObjects;
+
+namespace Microsoft.CodeAnalysis.Serialization
+{
+    internal readonly struct SolutionReplicationContext : IDisposable
+    {
+        private readonly ArrayBuilder<IDisposable> _resources;
+
+        private SolutionReplicationContext(ArrayBuilder<IDisposable> resources)
+            => _resources = resources;
+
+        public static SolutionReplicationContext Create()
+            => new(ArrayBuilder<IDisposable>.GetInstance());
+
+        public void AddResource(IDisposable resource)
+            => _resources.Add(resource);
+
+        public void Dispose()
+        {
+            // TODO: https://github.com/dotnet/roslyn/issues/49973
+            // Currently we don't dispose resources, only keep them alive.
+            // Shouldn't we dispose them? 
+            // _resources.All(resource => resource.Dispose());
+            _resources.Free();
+        }
+    }
+}

src/Workspaces/Core/Portable/Workspace/Solution/Checksum_Factory.cs

@@ -138,11 +138,12 @@ public static Checksum Create(WellKnownSynchronizationKind kind, ImmutableArray<
         public static Checksum Create<T>(WellKnownSynchronizationKind kind, T value, ISerializerService serializer)
         {
             using var stream = SerializableBytes.CreateWritableStream();
+            using var context = SolutionReplicationContext.Create();
 
             using (var objectWriter = new ObjectWriter(stream, leaveOpen: true))
             {
                 objectWriter.WriteInt32((int)kind);
-                serializer.Serialize(value, objectWriter, CancellationToken.None);
+                serializer.Serialize(value, objectWriter, context, CancellationToken.None);
             }
 
             stream.Position = 0;

src/Workspaces/CoreTestUtilities/Remote/TestSerializerService.cs

@@ -36,7 +36,7 @@ public TestSerializerService(HostWorkspaceServices workspaceServices)
         {
         }
 
-        public override void WriteMetadataReferenceTo(MetadataReference reference, ObjectWriter writer, CancellationToken cancellationToken)
+        public override void WriteMetadataReferenceTo(MetadataReference reference, ObjectWriter writer, SolutionReplicationContext context, CancellationToken cancellationToken)
         {
             var wellKnownReferenceName = s_wellKnownReferenceNames.GetValueOrDefault(reference, null);
             if (wellKnownReferenceName is not null)
@@ -47,7 +47,7 @@ public override void WriteMetadataReferenceTo(MetadataReference reference, Objec
             else
             {
                 writer.WriteBoolean(false);
-                base.WriteMetadataReferenceTo(reference, writer, cancellationToken);
+                base.WriteMetadataReferenceTo(reference, writer, context, cancellationToken);
             }
         }
 

src/Workspaces/Remote/Core/RemoteHostAssetSerialization.cs

@@ -35,7 +35,8 @@ public static async Task WriteDataAsync(ObjectWriter writer, SolutionAssetStorag
                 assetMap = await assetStorage.GetAssetsAsync(scopeId, checksums, cancellationToken).ConfigureAwait(false);
             }
 
-            WriteData(writer, singleAsset, assetMap, serializer, scopeId, checksums, cancellationToken);
+            var replicationContext = assetStorage.GetReplicationContext(scopeId);
+            WriteData(writer, singleAsset, assetMap, serializer, replicationContext, scopeId, checksums, cancellationToken);
         }
 
         public static void WriteData(
@@ -43,6 +44,7 @@ public static async Task WriteDataAsync(ObjectWriter writer, SolutionAssetStorag
             SolutionAsset? singleAsset,
             IReadOnlyDictionary<Checksum, SolutionAsset>? assetMap,
             ISerializerService serializer,
+            SolutionReplicationContext context,
             int scopeId,
             Checksum[] checksums,
             CancellationToken cancellationToken)
@@ -59,7 +61,7 @@ public static async Task WriteDataAsync(ObjectWriter writer, SolutionAssetStorag
             if (singleAsset != null)
             {
                 writer.WriteInt32(1);
-                WriteAsset(writer, serializer, checksums[0], singleAsset, cancellationToken);
+                WriteAsset(writer, serializer, context, checksums[0], singleAsset, cancellationToken);
                 return;
             }
 
@@ -68,10 +70,10 @@ public static async Task WriteDataAsync(ObjectWriter writer, SolutionAssetStorag
 
             foreach (var (checksum, asset) in assetMap)
             {
-                WriteAsset(writer, serializer, checksum, asset, cancellationToken);
+                WriteAsset(writer, serializer, context, checksum, asset, cancellationToken);
             }
 
-            static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Checksum checksum, SolutionAsset asset, CancellationToken cancellationToken)
+            static void WriteAsset(ObjectWriter writer, ISerializerService serializer, SolutionReplicationContext context, Checksum checksum, SolutionAsset asset, CancellationToken cancellationToken)
             {
                 checksum.WriteTo(writer);
                 writer.WriteInt32((int)asset.Kind);
@@ -79,7 +81,7 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check
                 // null is already indicated by checksum and kind above:
                 if (asset.Value is not null)
                 {
-                    serializer.Serialize(asset.Value, writer, cancellationToken);
+                    serializer.Serialize(asset.Value, writer, context, cancellationToken);
                 }
             }
         }

src/Workspaces/Remote/Core/SolutionAssetProvider.cs

@@ -20,6 +20,9 @@
 
 namespace Microsoft.CodeAnalysis.Remote
 {
+    /// <summary>
+    /// Provides solution assets present locally (in the current process) to a remote process where the solution is being replicated to.
+    /// </summary>
     internal sealed class SolutionAssetProvider : ISolutionAssetProvider
     {
         public const string ServiceName = ServiceDescriptors.ServiceNameTopLevelPrefix + ServiceDescriptors.ComponentName + ".SolutionAssetProvider";
@@ -37,6 +40,7 @@ public async ValueTask GetAssetsAsync(PipeWriter pipeWriter, int scopeId, Checks
         {
             var assetStorage = _services.GetRequiredService<ISolutionAssetStorageProvider>().AssetStorage;
             var serializer = _services.GetRequiredService<ISerializerService>();
+            var replicationContext = assetStorage.GetReplicationContext(scopeId);
 
             SolutionAsset? singleAsset = null;
             IReadOnlyDictionary<Checksum, SolutionAsset>? assetMap = null;
@@ -68,7 +72,7 @@ public async ValueTask GetAssetsAsync(PipeWriter pipeWriter, int scopeId, Checks
                 {
                     var stream = localPipe.Writer.AsStream(leaveOpen: false);
                     using var writer = new ObjectWriter(stream, leaveOpen: false, cancellationToken);
-                    RemoteHostAssetSerialization.WriteData(writer, singleAsset, assetMap, serializer, scopeId, checksums, cancellationToken);
+                    RemoteHostAssetSerialization.WriteData(writer, singleAsset, assetMap, serializer, replicationContext, scopeId, checksums, cancellationToken);
                 }
                 catch (Exception e) when (FatalError.ReportAndCatchUnlessCanceled(e, cancellationToken))
                 {

src/Workspaces/Remote/Core/SolutionAssetStorage.Scope.cs

@@ -21,7 +21,10 @@ public Scope(SolutionAssetStorage storages, PinnedSolutionInfo solutionInfo)
             }
 
             public void Dispose()
-                => Contract.ThrowIfFalse(_storages._solutionStates.TryRemove(SolutionInfo.ScopeId, out _));
+            {
+                Contract.ThrowIfFalse(_storages._solutionStates.TryRemove(SolutionInfo.ScopeId, out var entry));
+                entry.ReplicationContext.Dispose();
+            }
         }
     }
 }

src/Workspaces/Remote/Core/SolutionAssetStorage.cs

@@ -24,12 +24,10 @@ internal partial class SolutionAssetStorage
         /// <summary>
         /// Map from solution checksum scope id to its associated <see cref="SolutionState"/>.
         /// </summary>
-        private readonly ConcurrentDictionary<int, SolutionState> _solutionStates;
+        private readonly ConcurrentDictionary<int, (SolutionState Solution, SolutionReplicationContext ReplicationContext)> _solutionStates = new(concurrencyLevel: 2, capacity: 10);
 
-        public SolutionAssetStorage()
-        {
-            _solutionStates = new ConcurrentDictionary<int, SolutionState>(concurrencyLevel: 2, capacity: 10);
-        }
+        public SolutionReplicationContext GetReplicationContext(int scopeId)
+            => _solutionStates[scopeId].ReplicationContext;
 
         /// <summary>
         /// Adds given snapshot into the storage. This snapshot will be available within the returned <see cref="Scope"/>.
@@ -38,6 +36,7 @@ internal async ValueTask<Scope> StoreAssetsAsync(Solution solution, Cancellation
         {
             var solutionState = solution.State;
             var solutionChecksum = await solutionState.GetChecksumAsync(cancellationToken).ConfigureAwait(false);
+            var context = SolutionReplicationContext.Create();
 
             var id = Interlocked.Increment(ref s_scopeId);
             var solutionInfo = new PinnedSolutionInfo(
@@ -46,7 +45,7 @@ internal async ValueTask<Scope> StoreAssetsAsync(Solution solution, Cancellation
                 solutionState.WorkspaceVersion,
                 solutionChecksum);
 
-            Contract.ThrowIfFalse(_solutionStates.TryAdd(id, solutionState));
+            Contract.ThrowIfFalse(_solutionStates.TryAdd(id, (solutionState, context)));
 
             return new Scope(this, solutionInfo);
         }
@@ -62,7 +61,7 @@ internal async ValueTask<Scope> StoreAssetsAsync(Solution solution, Cancellation
                 return SolutionAsset.Null;
             }
 
-            var remotableData = await FindAssetAsync(_solutionStates[scopeId], checksum, cancellationToken).ConfigureAwait(false);
+            var remotableData = await FindAssetAsync(_solutionStates[scopeId].Solution, checksum, cancellationToken).ConfigureAwait(false);
             if (remotableData != null)
             {
                 return remotableData;
@@ -93,7 +92,7 @@ internal async ValueTask<Scope> StoreAssetsAsync(Solution solution, Cancellation
                 result[Checksum.Null] = SolutionAsset.Null;
             }
 
-            await FindAssetsAsync(_solutionStates[scopeId], checksumsToFind.Object, result, cancellationToken).ConfigureAwait(false);
+            await FindAssetsAsync(_solutionStates[scopeId].Solution, checksumsToFind.Object, result, cancellationToken).ConfigureAwait(false);
             if (result.Count == numberOfChecksumsToSearch)
             {
                 // no checksum left to find