Bicep.RpcClient - use PipeReader & PipeWriter APIs (#18212)

* Use PipeReader & PipeWriter APIs
* Ensure the child process is explicitly killed

Author: anthony-c-martin

src/Bicep.RpcClient.Tests/PublicApiTests.cs

@@ -41,6 +41,7 @@ public void Dependencies_should_be_minimal()
             // Be careful when adding new dependencies to the ClientTools assembly - this assembly is intentionally slim.
             // The assembly is used in Microsoft internal tools, where dependency management is complex, so we want to avoid transitively depending on ResourceStack
             "System.Collections.Immutable",
+            "System.IO.Pipelines",
             "System.Memory",
             "System.Text.Encodings.Web",
             "System.Text.Json",

src/Bicep.RpcClient/BicepClient.cs

@@ -4,8 +4,11 @@
 using System;
 using System.Diagnostics;
 using System.IO;
+using System.IO.Pipelines;
 using System.IO.Pipes;
 using System.Threading;
+using System.Threading.Tasks;
+using Bicep.RpcClient;
 using Bicep.RpcClient.JsonRpc;
 using Bicep.RpcClient.Models;
 
@@ -17,12 +20,13 @@ internal class BicepClient : IBicepClient
     private readonly JsonRpcClient jsonRpcClient;
     private readonly Task backgroundTask;
     private readonly CancellationTokenSource cts;
+    private string? cachedVersion;
 
     private BicepClient(NamedPipeServerStream pipeStream, Process cliProcess, JsonRpcClient jsonRpcClient, CancellationTokenSource cts)
     {
         this.cliProcess = cliProcess;
         this.jsonRpcClient = jsonRpcClient;
-        this.backgroundTask = CreateBackgroundTask(pipeStream, cliProcess, jsonRpcClient, cts.Token);
+        this.backgroundTask = jsonRpcClient.Listen(onComplete: pipeStream.Dispose, cts.Token);
         this.cts = cts;
     }
 
@@ -37,7 +41,7 @@ public static async Task<IBicepClient> Initialize(string bicepCliPath, Cancellat
         }
 
         var pipeName = Guid.NewGuid().ToString();
-        var pipeStream = new NamedPipeServerStream(pipeName, PipeDirection.InOut, NamedPipeServerStream.MaxAllowedServerInstances, PipeTransmissionMode.Byte, PipeOptions.Asynchronous);
+        var pipeStream = new NamedPipeServerStream(pipeName, PipeDirection.InOut, 1, PipeTransmissionMode.Byte, System.IO.Pipes.PipeOptions.Asynchronous);
 
         var psi = new ProcessStartInfo
         {
@@ -54,29 +58,11 @@ public static async Task<IBicepClient> Initialize(string bicepCliPath, Cancellat
 
         await pipeStream.WaitForConnectionAsync(cancellationToken).ConfigureAwait(false);
         var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        var client = new JsonRpcClient(pipeStream, pipeStream);
+        var client = new JsonRpcClient(PipeReader.Create(pipeStream), PipeWriter.Create(pipeStream));
 
         return new BicepClient(pipeStream, cliProcess, client, cts);
     }
 
-    private static Task CreateBackgroundTask(NamedPipeServerStream pipeStream, Process cliProcess, JsonRpcClient jsonRpcClient, CancellationToken cancellationToken)
-        => Task.Run(async () =>
-        {
-            try
-            {
-                await jsonRpcClient.ReadLoop(cancellationToken).ConfigureAwait(false);
-            }
-            catch (OperationCanceledException)
-            {
-                // Expected when disposing
-            }
-            finally
-            {
-                pipeStream.Dispose();
-                cliProcess.Dispose();
-            }
-        }, cancellationToken);
-
     /// <inheritdoc/>
     public Task<CompileResponse> Compile(CompileRequest request, CancellationToken cancellationToken)
         => jsonRpcClient.SendRequest<CompileRequest, CompileResponse>("bicep/compile", request, cancellationToken);
@@ -114,8 +100,15 @@ public async Task<GetSnapshotResponse> GetSnapshot(GetSnapshotRequest request, C
     /// <inheritdoc/>
     public async Task<string> GetVersion(CancellationToken cancellationToken)
     {
-        var response = await jsonRpcClient.SendRequest<VersionRequest, VersionResponse>("bicep/version", new(), cancellationToken).ConfigureAwait(false);
-        return response.Version;
+        // This method is called frequently for version checks - cache the result
+        if (cachedVersion is null)
+        {
+            var response = await jsonRpcClient.SendRequest<VersionRequest, VersionResponse>("bicep/version", new(), cancellationToken).ConfigureAwait(false);
+            // No locks needed here, since replacing a reference is atomic
+            cachedVersion = response.Version;
+        }
+
+        return cachedVersion;
     }
 
     private async Task EnsureMinimumVersion(string requiredVersion, string operationName, CancellationToken cancellationToken)
@@ -130,9 +123,15 @@ private async Task EnsureMinimumVersion(string requiredVersion, string operation
     public void Dispose()
     {
         cts.Cancel();
-        if (!cliProcess.HasExited)
+        jsonRpcClient.Dispose();
+        try
+        {
+            cliProcess.Kill();
+            // wait for 10 seconds for the process to exit
+            cliProcess.WaitForExit(10000);
+        }
+        catch (InvalidOperationException)
         {
-            cliProcess.WaitForExit();
         }
     }
 }

src/Bicep.RpcClient/JsonRpc/JsonRpcClient.cs

@@ -1,17 +1,23 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System;
+using System.Buffers;
 using System.Collections.Concurrent;
-using System.Net.Security;
+using System.IO;
+using System.IO.Pipelines;
+using System.Linq;
 using System.Text;
 using System.Text.Encodings.Web;
 using System.Text.Json;
 using System.Text.Json.Nodes;
 using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Bicep.RpcClient.JsonRpc;
 
-internal class JsonRpcClient(Stream reader, Stream writer) : IDisposable
+internal class JsonRpcClient(PipeReader reader, PipeWriter writer) : IDisposable
 {
     private record JsonRpcRequest<T>(
         string Jsonrpc,
@@ -20,7 +26,6 @@ private record JsonRpcRequest<T>(
         int Id);
 
     private record MinimalJsonRpcResponse(
-        string Jsonrpc,
         int Id);
 
     private record JsonRpcResponse<T>(
@@ -34,10 +39,9 @@ private record JsonRpcError(
         string Message,
         JsonNode? Data);
 
-    private readonly byte[] terminator = "\r\n\r\n"u8.ToArray();
     private int nextId = 0;
     private readonly SemaphoreSlim writeSemaphore = new(1, 1);
-    private readonly ConcurrentDictionary<int, TaskCompletionSource<string>> pendingResponses = new();
+    private readonly ConcurrentDictionary<int, TaskCompletionSource<byte[]>> pendingResponses = new();
 
     private readonly JsonSerializerOptions jsonSerializerOptions = new()
     {
@@ -55,20 +59,20 @@ public async Task<TResponse> SendRequest<TRequest, TResponse>(string method, TRe
         var requestContent = JsonSerializer.Serialize(jsonRpcRequest, jsonSerializerOptions);
         var requestLength = Encoding.UTF8.GetByteCount(requestContent);
         var rawRequest = $"Content-Length: {requestLength}\r\n\r\n{requestContent}";
-        var requestBytes = Encoding.UTF8.GetBytes(rawRequest);
 
         await writeSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
         try
         {
-            await writer.WriteAsync(requestBytes, 0, requestBytes.Length, cancellationToken).ConfigureAwait(false);
+            var requestBytes = Encoding.UTF8.GetBytes(rawRequest).AsMemory();
+            await writer.WriteAsync(requestBytes, cancellationToken).ConfigureAwait(false);
             await writer.FlushAsync(cancellationToken).ConfigureAwait(false);
         }
         finally
         {
             writeSemaphore.Release();
         }
 
-        var tcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var tcs = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
         if (!pendingResponses.TryAdd(currentId, tcs))
         {
             throw new InvalidOperationException($"A request with ID {currentId} is already pending.");
@@ -87,13 +91,34 @@ public async Task<TResponse> SendRequest<TRequest, TResponse>(string method, TRe
         return jsonRpcResponse.Result;
     }
 
-    public async Task ReadLoop(CancellationToken cancellationToken)
+    public Task Listen(Action onComplete, CancellationToken cancellationToken)
+        => Task.Run(async () =>
+        {
+            try
+            {
+                await ListenInternal(cancellationToken).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                // Expected when disposing
+            }
+            finally
+            {
+                onComplete();
+            }
+        }, cancellationToken);
+
+    private async Task ListenInternal(CancellationToken cancellationToken)
     {
         while (true)
         {
             try
             {
                 var message = await ReadMessage(cancellationToken).ConfigureAwait(false);
+                if (message is null)
+                {
+                    return;
+                }
 
                 var response = JsonSerializer.Deserialize<MinimalJsonRpcResponse>(message, jsonSerializerOptions)
                     ?? throw new InvalidOperationException("Failed to deserialize JSON-RPC response");
@@ -105,77 +130,134 @@ public async Task ReadLoop(CancellationToken cancellationToken)
             }
             catch (Exception) when (cancellationToken.IsCancellationRequested)
             {
-                reader.Dispose();
-                writer.Dispose();
+                await reader.CompleteAsync().ConfigureAwait(false);
+                await writer.CompleteAsync().ConfigureAwait(false);
                 break;
             }
         }
     }
 
-    private async Task<string> ReadUntilTerminator(CancellationToken cancellationToken)
-    {
-        using var outputStream = new MemoryStream();
-        var patternIndex = 0;
-        var byteBuffer = new byte[1];
+    private record Headers(
+        int ContentLength);
 
+    private async Task<Headers?> ReadHeaders(CancellationToken cancellationToken)
+    {
+        int? contentLength = null;
         while (true)
         {
-            await ReadExactly(byteBuffer, byteBuffer.Length, cancellationToken).ConfigureAwait(false);
+            var readResult = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
 
-            await outputStream.WriteAsync(byteBuffer, 0, byteBuffer.Length, cancellationToken).ConfigureAwait(false);
-            patternIndex = terminator[patternIndex] == byteBuffer[0] ? patternIndex + 1 : 0;
-            if (patternIndex == terminator.Length)
+            if (readResult.Buffer.IsEmpty && readResult.IsCompleted)
             {
-                outputStream.Position = 0;
-                outputStream.SetLength(outputStream.Length - terminator.Length);
-                // return stream as string
-                return Encoding.UTF8.GetString(outputStream.ToArray());
+                return null; // remote end disconnected at a reasonable place.
+            }
+
+            var lf = readResult.Buffer.PositionOf((byte)'\n');
+            if (!lf.HasValue)
+            {
+                if (readResult.IsCompleted)
+                {
+                    throw new EndOfStreamException();
+                }
+
+                // Indicate that we can't find what we're looking for and read again.
+                reader.AdvanceTo(readResult.Buffer.Start, readResult.Buffer.End);
+                continue;
+            }
+
+            var line = readResult.Buffer.Slice(0, lf.Value);
+
+            // Verify the line ends with an \r (that precedes the \n we already found)
+            var cr = line.PositionOf((byte)'\r');
+            if (!cr.HasValue || !line.GetPosition(1, cr.Value).Equals(lf))
+            {
+                throw new InvalidOperationException("Header does not end with expected \r\n character sequence");
+            }
+
+            // Trim off the \r now that we confirmed it was there.
+            line = line.Slice(0, line.Length - 1);
+
+            if (line.Length > 0)
+            {
+                var lineText = Encoding.UTF8.GetString(line.ToArray());
+                var split = lineText.Split([':'], 2);
+                if (split.Length != 2)
+                {
+                    throw new InvalidOperationException("Colon not found in header.");
+                }
+
+                var headerName = split[0].Trim();
+                var headerValue = split[1].Trim();
+
+                if (headerName == "Content-Length")
+                {
+                    contentLength = int.Parse(headerValue);
+                }
+            }
+
+            // Advance to the next line.
+            reader.AdvanceTo(readResult.Buffer.GetPosition(1, lf.Value));
+
+            if (line.Length == 0)
+            {
+                // We found the empty line that constitutes the end of the HTTP headers.
+                break;
             }
         }
-    }
 
-    private async Task<string> ReadContent(int length, CancellationToken cancellationToken)
-    {
-        var byteBuffer = new byte[length];
-        await ReadExactly(byteBuffer, length, cancellationToken).ConfigureAwait(false);
+        if (!contentLength.HasValue)
+        {
+            throw new InvalidOperationException("Failed to obtain Content-Length header");
+        }
 
-        return Encoding.UTF8.GetString(byteBuffer);
+        return new(contentLength.Value);
     }
 
-    private async Task<string> ReadMessage(CancellationToken cancellationToken)
+    protected async ValueTask<ReadResult> ReadAtLeastAsync(int requiredBytes, bool allowEmpty, CancellationToken cancellationToken)
     {
-        var header = await ReadUntilTerminator(cancellationToken).ConfigureAwait(false);
-        var parsed = header.Split(':').Select(x => x.Trim()).ToArray();
+        var readResult = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
+        while (readResult.Buffer.Length < requiredBytes && !readResult.IsCompleted && !readResult.IsCanceled)
+        {
+            reader.AdvanceTo(readResult.Buffer.Start, readResult.Buffer.End);
+            readResult = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
+        }
 
-        if (parsed.Length != 2 ||
-            !parsed[0].Equals("Content-Length", StringComparison.OrdinalIgnoreCase) ||
-            !int.TryParse(parsed[1], out var contentLength) ||
-            contentLength <= 0)
+        if (allowEmpty && readResult.Buffer.Length == 0)
         {
-            throw new InvalidOperationException($"Invalid header: {header}");
+            return readResult;
         }
 
-        return await ReadContent(contentLength, cancellationToken).ConfigureAwait(false);
+        if (readResult.Buffer.Length < requiredBytes)
+        {
+            throw readResult.IsCompleted ? new EndOfStreamException() :
+                readResult.IsCanceled ? new OperationCanceledException() :
+                throw new InvalidOperationException(); // should be unreachable
+        }
+
+        return readResult;
     }
 
-    private async ValueTask ReadExactly(byte[] buffer, int minimumBytes, CancellationToken cancellationToken)
+    private async Task<byte[]?> ReadMessage(CancellationToken cancellationToken)
     {
-        var totalRead = 0;
-        while (totalRead < minimumBytes)
+        var headers = await ReadHeaders(cancellationToken).ConfigureAwait(false);
+        if (headers is null)
         {
-            var read = await reader.ReadAsync(buffer, totalRead, minimumBytes - totalRead, cancellationToken).ConfigureAwait(false);
-            if (read == 0)
-            {
-                throw new EndOfStreamException("Stream closed before reading expected number of bytes");
-            }
-
-            totalRead += read;
+            return null;
         }
+
+        var readResult = await ReadAtLeastAsync(headers.ContentLength, allowEmpty: false, cancellationToken).ConfigureAwait(false);
+
+        var contentBuffer = readResult.Buffer.Slice(0, headers.ContentLength);
+        var output = contentBuffer.ToArray();
+
+        reader.AdvanceTo(contentBuffer.End);
+
+        return output;
     }
 
     public void Dispose()
     {
-        writer.Dispose();
-        reader.Dispose();
+        writer.Complete();
+        reader.Complete();
     }
 }