From 6abfe026f7fe29f4ed88b1b0e4b4e114dc1bb885 Mon Sep 17 00:00:00 2001
From: Andrew Arnott <andrew.arnott@microsoft.com>
Date: Tue, 19 Sep 2023 17:54:28 -0600
Subject: [PATCH] Support `SystemTextJsonFormatter` in
 `NewLineDelimitedMessageHandler`

---
 .../NewLineDelimitedMessageHandler.cs         | 30 ++++++++++++++-----
 .../NewLineDelimitedMessageHandlerTests.cs    | 13 ++++----
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/src/StreamJsonRpc/NewLineDelimitedMessageHandler.cs b/src/StreamJsonRpc/NewLineDelimitedMessageHandler.cs
index 29ea48df..a7ceceb8 100644
--- a/src/StreamJsonRpc/NewLineDelimitedMessageHandler.cs
+++ b/src/StreamJsonRpc/NewLineDelimitedMessageHandler.cs
@@ -5,6 +5,7 @@
 using System.Globalization;
 using System.IO.Pipelines;
 using System.Text;
+using Nerdbank.Streams;
 using StreamJsonRpc.Protocol;
 using StreamJsonRpc.Reflection;
 
@@ -19,6 +20,11 @@ namespace StreamJsonRpc;
 /// </remarks>
 public class NewLineDelimitedMessageHandler : PipeMessageHandler
 {
+    /// <summary>
+    /// The <see cref="IBufferWriter{T}"/> that buffers an outgoing message.
+    /// </summary>
+    private readonly Sequence<byte> contentSequenceBuilder = new Sequence<byte>(ArrayPool<byte>.Shared);
+
     /// <summary>
     /// Backing field for the <see cref="NewLine"/> property.
     /// </summary>
@@ -105,7 +111,22 @@ protected override void Write(JsonRpcMessage content, CancellationToken cancella
         Assumes.NotNull(this.Writer);
 
         cancellationToken.ThrowIfCancellationRequested();
-        this.Formatter.Serialize(this.Writer, content);
+
+        // Some formatters (e.g. MessagePackFormatter) needs the encoded form in order to produce JSON for tracing.
+        // Other formatters (e.g. JsonMessageFormatter) would prefer to do its own tracing while it still has a JToken.
+        // We only help the formatters that need the byte-encoded form here. The rest can do it themselves.
+        if (this.Formatter is IJsonRpcFormatterTracingCallbacks tracer)
+        {
+            this.Formatter.Serialize(this.contentSequenceBuilder, content);
+            tracer.OnSerializationComplete(content, this.contentSequenceBuilder);
+            this.Writer.Write(this.contentSequenceBuilder);
+            this.contentSequenceBuilder.Reset();
+        }
+        else
+        {
+            this.Formatter.Serialize(this.Writer, content);
+        }
+
         this.Writer.Write(this.newLineBytes.Span);
     }
 
@@ -180,13 +201,6 @@ private static ReadOnlyMemory<byte> GetLineFeedSequence(Encoding encoding, NewLi
     /// </summary>
     private void CommonConstructor()
     {
-        if (this.Formatter is IJsonRpcFormatterTracingCallbacks)
-        {
-            // Such a formatter requires that we make their own written bytes available back to them for logging purposes.
-            // We haven't implemented support for such, and the JsonMessageFormatter doesn't need it, so no need to.
-            throw new NotSupportedException("Formatters that implement " + nameof(IJsonRpcFormatterTracingCallbacks) + " are not supported. Try using " + nameof(JsonMessageFormatter) + ".");
-        }
-
         if (this.Formatter.Encoding.WebName != Encoding.UTF8.WebName)
         {
             throw new NotSupportedException("Only UTF-8 formatters are supported.");
diff --git a/test/StreamJsonRpc.Tests/NewLineDelimitedMessageHandlerTests.cs b/test/StreamJsonRpc.Tests/NewLineDelimitedMessageHandlerTests.cs
index 7ca5653d..782323bf 100644
--- a/test/StreamJsonRpc.Tests/NewLineDelimitedMessageHandlerTests.cs
+++ b/test/StreamJsonRpc.Tests/NewLineDelimitedMessageHandlerTests.cs
@@ -53,11 +53,11 @@ public void NewLine()
         Assert.Equal(NewLineDelimitedMessageHandler.NewLineStyle.Lf, handler.NewLine);
     }
 
-    [Fact]
-    public async Task Reading_MixedLineEndings()
+    [Theory, PairwiseData]
+    public async Task Reading_MixedLineEndings(bool useSystemTextJson)
     {
         var pipe = new Pipe();
-        var handler = new NewLineDelimitedMessageHandler(null, pipe.Reader, new JsonMessageFormatter());
+        var handler = new NewLineDelimitedMessageHandler(null, pipe.Reader, useSystemTextJson ? new SystemTextJsonFormatter() : new JsonMessageFormatter());
 
         // Send messages with mixed line endings to the handler..
         var testFormatter = new JsonMessageFormatter();
@@ -105,11 +105,11 @@ public async Task Reading_IncompleteLine()
         Assert.True(msg is JsonRpcRequest);
     }
 
-    [Fact]
-    public async Task Writing_MixedLineEndings()
+    [Theory, PairwiseData]
+    public async Task Writing_MixedLineEndings(bool useSystemTextJson)
     {
         var pipe = new Pipe();
-        var handler = new NewLineDelimitedMessageHandler(pipe.Writer, null, new JsonMessageFormatter());
+        var handler = new NewLineDelimitedMessageHandler(pipe.Writer, null, useSystemTextJson ? new SystemTextJsonFormatter() : new JsonMessageFormatter());
 
         // Use the handler to write out a couple messages with mixed line endings.
         await handler.WriteAsync(this.mockMessages[0], this.TimeoutToken); // CRLF
@@ -119,6 +119,7 @@ public async Task Writing_MixedLineEndings()
 
         using var streamReader = new StreamReader(pipe.Reader.AsStream(), handler.Formatter.Encoding);
         string allMessages = await streamReader.ReadToEndAsync();
+        this.Logger.WriteLine(allMessages);
 
         // Use CR and LF counts to quickly figure whether our new line styles were honored.
         Assert.Equal(3, allMessages.Split('\n').Length);