Change all messages to pass back their checksum. Switch to callback style

src/Workspaces/CoreTestUtilities/Fakes/SimpleAssetSource.cs

@@ -19,9 +19,8 @@ namespace Microsoft.CodeAnalysis.Remote.Testing;
 internal sealed class SimpleAssetSource(ISerializerService serializerService, IReadOnlyDictionary<Checksum, object> map) : IAssetSource
 {
     public ValueTask GetAssetsAsync<T, TArg>(
-        Checksum solutionChecksum, AssetPath assetPath, ReadOnlyMemory<Checksum> checksums, ISerializerService deserializerService, Action<int, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
+        Checksum solutionChecksum, AssetPath assetPath, ReadOnlyMemory<Checksum> checksums, ISerializerService deserializerService, Action<Checksum, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
     {
-        var index = 0;
         foreach (var checksum in checksums.Span)
         {
             Contract.ThrowIfFalse(map.TryGetValue(checksum, out var data));
@@ -38,8 +37,7 @@ public ValueTask GetAssetsAsync<T, TArg>(
             using var reader = ObjectReader.GetReader(stream, leaveOpen: true, cancellationToken);
             var asset = deserializerService.Deserialize(data.GetWellKnownSynchronizationKind(), reader, cancellationToken);
             Contract.ThrowIfNull(asset);
-            callback(index, (T)asset, arg);
-            index++;
+            callback(checksum, (T)asset, arg);
         }
 
         return ValueTaskFactory.CompletedTask;

src/Workspaces/Remote/Core/RemoteHostAssetSerialization.cs

@@ -39,33 +39,25 @@ public static async ValueTask WriteDataAsync(
 
             Debug.Assert(assetMap != null);
 
-            for (var i = 0; i < checksums.Span.Length; i++)
+            foreach (var (checksum, asset) in assetMap)
             {
-                var checksum = checksums.Span[i];
-                var asset = assetMap[checksum];
+                Contract.ThrowIfNull(asset);
+
+                var kind = asset.GetWellKnownSynchronizationKind();
+                checksum.WriteTo(writer);
+                writer.WriteInt32((int)kind);
+                serializer.Serialize(asset, writer, context, cancellationToken);
 
                 // We flush after each item as that forms a reasonably sized chunk of data to want to then send over the
                 // pipe for the reader on the other side to read.  This allows the item-writing to remain entirely
                 // synchronous without any blocking on async flushing, while also ensuring that we're not buffering the
                 // entire stream of data into the pipe before it gets sent to the other side.
-                WriteAsset(writer, serializer, context, asset, cancellationToken);
                 await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
             }
-
-            return;
-
-            static void WriteAsset(ObjectWriter writer, ISerializerService serializer, SolutionReplicationContext context, object asset, CancellationToken cancellationToken)
-            {
-                Contract.ThrowIfNull(asset);
-                var kind = asset.GetWellKnownSynchronizationKind();
-                writer.WriteInt32((int)kind);
-
-                serializer.Serialize(asset, writer, context, cancellationToken);
-            }
         }
 
         public static ValueTask ReadDataAsync<T, TArg>(
-            PipeReader pipeReader, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<int, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
+            PipeReader pipeReader, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<Checksum, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
         {
             // Suppress ExecutionContext flow for asynchronous operations operate on the pipe. In addition to avoiding
             // ExecutionContext allocations, this clears the LogicalCallContext and avoids the need to clone data set by
@@ -77,15 +69,15 @@ public static ValueTask ReadDataAsync<T, TArg>(
             return ReadDataSuppressedFlowAsync(pipeReader, solutionChecksum, objectCount, serializerService, callback, arg, cancellationToken);
 
             static async ValueTask ReadDataSuppressedFlowAsync(
-                PipeReader pipeReader, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<int, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
+                PipeReader pipeReader, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<Checksum, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
             {
                 using var stream = await pipeReader.AsPrebufferedStreamAsync(cancellationToken).ConfigureAwait(false);
                 ReadData(stream, solutionChecksum, objectCount, serializerService, callback, arg, cancellationToken);
             }
         }
 
         public static void ReadData<T, TArg>(
-            Stream stream, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<int, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
+            Stream stream, Checksum solutionChecksum, int objectCount, ISerializerService serializerService, Action<Checksum, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
         {
             using var reader = ObjectReader.GetReader(stream, leaveOpen: true, cancellationToken);
 
@@ -96,12 +88,13 @@ public static void ReadData<T, TArg>(
 
             for (int i = 0, n = objectCount; i < n; i++)
             {
+                var checksum = Checksum.ReadFrom(reader);
                 var kind = (WellKnownSynchronizationKind)reader.ReadInt32();
 
                 // in service hub, cancellation means simply closed stream
                 var result = serializerService.Deserialize(kind, reader, cancellationToken);
                 Contract.ThrowIfNull(result);
-                callback(i, (T)result, arg);
+                callback(checksum, (T)result, arg);
             }
         }
     }

src/Workspaces/Remote/ServiceHub/Host/AssetProvider.cs

@@ -284,18 +284,26 @@ private async ValueTask SynchronizeAssetsAsync<T, TArg>(
                 if (missingChecksumsCount > 0)
                 {
                     var missingChecksumsMemory = new ReadOnlyMemory<Checksum>(missingChecksums, 0, missingChecksumsCount);
+                    Contract.ThrowIfTrue(missingChecksumsMemory.Length == 0);
 
-                    await RequestAssetsAsync(
-                        assetPath, missingChecksumsMemory,
+#if NETCOREAPP
+                    Contract.ThrowIfTrue(missingChecksumsMemory.Span.Contains(Checksum.Null));
+#else
+                    Contract.ThrowIfTrue(missingChecksumsMemory.Span.IndexOf(Checksum.Null) >= 0);
+#endif
+
+                    await _assetSource.GetAssetsAsync(
+                        _solutionChecksum, assetPath, missingChecksumsMemory, _serializerService,
                         static (
-                            int index,
+                            Checksum missingChecksum,
                             T missingAsset,
                             (AssetProvider assetProvider, Checksum[] missingChecksums, Action<Checksum, T, TArg>? callback, TArg? arg) tuple) =>
                         {
-                            var missingChecksum = tuple.missingChecksums[index];
+                            // Add to cache, returning the existing asset if it was added by another thread.
+                            missingAsset = (T)tuple.assetProvider._assetCache.GetOrAdd(missingChecksum, missingAsset!);
 
+                            // Let our caller know about the asset if they're asking for it.
                             tuple.callback?.Invoke(missingChecksum, missingAsset, tuple.arg!);
-                            tuple.assetProvider._assetCache.GetOrAdd(missingChecksum, missingAsset!);
                         },
                         (this, missingChecksums, callback, arg),
                         cancellationToken).ConfigureAwait(false);
@@ -307,22 +315,5 @@ await RequestAssetsAsync(
                     s_checksumPool.Free(missingChecksums);
             }
         }
-
-        return;
-    }
-
-    private async ValueTask RequestAssetsAsync<T, TArg>(
-        AssetPath assetPath, ReadOnlyMemory<Checksum> checksums, Action<int, T, TArg> callback, TArg arg, CancellationToken cancellationToken)
-    {
-#if NETCOREAPP
-        Contract.ThrowIfTrue(checksums.Span.Contains(Checksum.Null));
-#else
-        Contract.ThrowIfTrue(checksums.Span.IndexOf(Checksum.Null) >= 0);
-#endif
-
-        if (checksums.Length == 0)
-            return;
-
-        await _assetSource.GetAssetsAsync(_solutionChecksum, assetPath, checksums, _serializerService, callback, arg, cancellationToken).ConfigureAwait(false);
     }
 }

src/Workspaces/Remote/ServiceHub/Host/IAssetSource.cs

@@ -20,7 +20,7 @@ ValueTask GetAssetsAsync<T, TArg>(
         AssetPath assetPath,
         ReadOnlyMemory<Checksum> checksums,
         ISerializerService serializerService,
-        Action<int, T, TArg> callback,
+        Action<Checksum, T, TArg> callback,
         TArg arg,
         CancellationToken cancellationToken);
 }

src/Workspaces/Remote/ServiceHub/Host/SolutionAssetSource.cs

@@ -21,7 +21,7 @@ public async ValueTask GetAssetsAsync<T, TArg>(
         AssetPath assetPath,
         ReadOnlyMemory<Checksum> checksums,
         ISerializerService serializerService,
-        Action<int, T, TArg> assetCallback,
+        Action<Checksum, T, TArg> assetCallback,
         TArg arg,
         CancellationToken cancellationToken)
     {

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/SerializableBytes.cs

@@ -6,10 +6,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
-using System.Diagnostics.CodeAnalysis;
 using System.IO;
-using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -95,7 +92,7 @@ private static void BlowChunks(byte[][]? chunks)
     internal static PooledStream CreateWritableStream()
         => new ReadWriteStream();
 
-    public class PooledStream : Stream
+    public abstract class PooledStream : Stream
     {
         protected List<byte[]> chunks;
 
@@ -109,13 +106,7 @@ protected PooledStream(long length, List<byte[]> chunks)
             this.chunks = chunks;
         }
 
-        public override long Length
-        {
-            get
-            {
-                return this.length;
-            }
-        }
+        public override long Length => this.length;
 
         public override bool CanRead => true;
 
@@ -130,10 +121,7 @@ public override void Flush()
 
         public override long Position
         {
-            get
-            {
-                return this.position;
-            }
+            get => this.position;
 
             set
             {
@@ -176,16 +164,9 @@ public override long Seek(long offset, SeekOrigin origin)
         public override int ReadByte()
         {
             if (position >= length)
-            {
                 return -1;
-            }
-
-            var currentIndex = CurrentChunkIndex;
-            var chunk = chunks[currentIndex];
-
-            var currentOffset = CurrentChunkOffset;
-            var result = chunk[currentOffset];
 
+            var result = chunks[CurrentChunkIndex][CurrentChunkOffset];
             this.position++;
             return result;
         }
@@ -295,7 +276,7 @@ public override void Write(byte[] buffer, int offset, int count)
     {
     }
 
-    private class ReadWriteStream : PooledStream
+    private sealed class ReadWriteStream : PooledStream
     {
         public ReadWriteStream()
             : base(length: 0, chunks: SharedPools.BigDefault<List<byte[]>>().AllocateAndClear())