diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/ChunkedPipeBufferWriter.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/ChunkedPipeBufferWriter.cs new file mode 100644 index 00000000000..62ed2a349fb --- /dev/null +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/ChunkedPipeBufferWriter.cs @@ -0,0 +1,200 @@ +using System.Buffers; +using System.Diagnostics.CodeAnalysis; +using System.IO.Pipelines; + +namespace HotChocolate.Transport.Formatters; + +/// +/// A pipe writer that buffers small writes and copies them to the underlying pipe when the +/// buffer is full or a flush is requested. A write of bytes +/// or more goes to the pipe directly. The buffer is held only between the first buffered write +/// and the next flush. It is not thread safe. +/// +internal sealed class ChunkedPipeBufferWriter : PipeWriter, IDisposable +{ + private const int ChunkSize = 64 * 1024; + + /// + /// The size hint from which a write bypasses the buffer and goes to the pipe. + /// + internal const int DirectWriteThreshold = 4 * 1024; + + [ThreadStatic] + private static ChunkedPipeBufferWriter? s_cached; + + private PipeWriter _pipe = default!; + private byte[] _buffer = []; + private int _position; + private bool _writeToPipe; + private bool _rented; + + private ChunkedPipeBufferWriter() + { + } + + /// + /// Gets a writer for . The writer must be disposed after its last flush. + /// + public static ChunkedPipeBufferWriter Rent(PipeWriter pipe) + { + var writer = s_cached ?? new ChunkedPipeBufferWriter(); + s_cached = null; + writer._pipe = pipe; + writer._rented = true; + return writer; + } + + /// + /// Gets a value indicating whether the writer currently holds a buffer. + /// + internal bool HasBuffer => _buffer.Length > 0; + + public override void Advance(int bytes) + { + if (_writeToPipe) + { + _writeToPipe = false; + _pipe.Advance(bytes); + return; + } + + if ((uint)bytes > (uint)(_buffer.Length - _position)) + { + ThrowAdvanceOutOfRange(bytes); + } + + _position += bytes; + } + + public override Memory GetMemory(int sizeHint = 0) + { + if (EnsureCapacity(sizeHint)) + { + return _buffer.AsMemory(_position); + } + + return _pipe.GetMemory(sizeHint); + } + + public override Span GetSpan(int sizeHint = 0) + { + if (EnsureCapacity(sizeHint)) + { + return _buffer.AsSpan(_position); + } + + return _pipe.GetSpan(sizeHint); + } + + /// + /// Copies the buffered bytes to the pipe, releases the buffer and flushes the pipe. + /// + public override ValueTask FlushAsync(CancellationToken cancellationToken = default) + { + CopyToPipe(); + ReleaseBuffer(); + return _pipe.FlushAsync(cancellationToken); + } + + public override void CancelPendingFlush() => _pipe.CancelPendingFlush(); + + public override void Complete(Exception? exception = null) + { + if (exception is null) + { + CopyToPipe(); + } + + ReleaseBuffer(); + _pipe.Complete(exception); + } + + /// + /// Discards the buffered bytes that have not been copied to the pipe yet. + /// Bytes that were already copied to the pipe are kept. + /// + public void Discard() + { + _writeToPipe = false; + ReleaseBuffer(); + } + + private void CopyToPipe() + { + if (_position > 0) + { + _pipe.Write(_buffer.AsSpan(0, _position)); + _position = 0; + } + } + + private void ReleaseBuffer() + { + if (_buffer.Length > 0) + { + ArrayPool.Shared.Return(_buffer); + _buffer = []; + _position = 0; + } + } + + /// + /// Returns true when the buffer serves bytes and + /// false when the next write and its go to the pipe. + /// + private bool EnsureCapacity(int sizeHint) + { + _writeToPipe = false; + + if (sizeHint >= DirectWriteThreshold) + { + CopyToPipe(); + _writeToPipe = true; + return false; + } + + if (sizeHint < 1) + { + sizeHint = 1; + } + + if (_buffer.Length == 0) + { + _buffer = ArrayPool.Shared.Rent(ChunkSize); + } + else if (_buffer.Length - _position < sizeHint) + { + CopyToPipe(); + } + + return true; + } + + [DoesNotReturn] + private static void ThrowAdvanceOutOfRange(int bytes) + { + if (bytes < 0) + { + throw ThrowHelper.ChunkedPipeBufferWriter_AdvanceNegative(bytes); + } + + throw ThrowHelper.ChunkedPipeBufferWriter_AdvancePastBuffer(); + } + + /// + /// Releases the writer for reuse. Bytes that were not flushed are discarded. + /// + public void Dispose() + { + if (!_rented) + { + return; + } + + ReleaseBuffer(); + _pipe = default!; + _writeToPipe = false; + _rented = false; + s_cached ??= this; + } +} diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/EventStreamResultFormatter.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/EventStreamResultFormatter.cs index 141835a757d..c09d11cacaa 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Formatters/EventStreamResultFormatter.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/EventStreamResultFormatter.cs @@ -51,10 +51,11 @@ IResponseStream responseStream private async ValueTask FormatOperationResultAsync( OperationResult operationResult, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); OperationResultFormatterContext? formatContext = null; var scope = Log.FormatOperationResultStart(); @@ -83,10 +84,11 @@ private async ValueTask FormatOperationResultAsync( private async ValueTask FormatResultBatchAsync( OperationResultBatch resultBatch, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); Exception? exception = null; using var semaphore = new SemaphoreSlim(1, 1); List? streams = null; @@ -125,6 +127,7 @@ private async ValueTask FormatResultBatchAsync( catch (Exception ex) { scope?.AddError(ex); + writer.Discard(); } finally { @@ -136,7 +139,7 @@ private async ValueTask FormatResultBatchAsync( } case IResponseStream responseStream: - keepAlive ??= new KeepAliveJob(semaphore, writer); + keepAlive ??= new KeepAliveJob(semaphore, pipe); var formatter = new StreamFormatter( _payloadFormatter, useIncrementalRfc1, @@ -186,16 +189,17 @@ private async ValueTask FormatResultBatchAsync( private async ValueTask FormatResponseStreamAsync( IResponseStream responseStream, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); Exception? exception = null; using var semaphore = new SemaphoreSlim(1, 1); try { - using (var keepAlive = new KeepAliveJob(semaphore, writer)) + using (var keepAlive = new KeepAliveJob(semaphore, pipe)) { var formatter = new StreamFormatter( _payloadFormatter, @@ -280,7 +284,7 @@ private sealed class StreamFormatter( KeepAliveJob keepAliveJob, IResponseStream responseStream, SemaphoreSlim semaphore, - PipeWriter writer) + ChunkedPipeBufferWriter writer) { private OperationResultFormatterContext? _formatContext; @@ -309,6 +313,7 @@ public async Task ProcessAsync(CancellationToken ct) catch (Exception ex) { scope?.AddError(ex); + writer.Discard(); throw; } finally diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/HotChocolate.Transport.Formatters.csproj b/src/HotChocolate/AspNetCore/src/Transport.Formatters/HotChocolate.Transport.Formatters.csproj index a27e0250b21..01c1d41381e 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Formatters/HotChocolate.Transport.Formatters.csproj +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/HotChocolate.Transport.Formatters.csproj @@ -11,6 +11,10 @@ true + + + + diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonLinesResultFormatter.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonLinesResultFormatter.cs index 9234eaeb545..c0454493b3a 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonLinesResultFormatter.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonLinesResultFormatter.cs @@ -43,10 +43,11 @@ IResponseStream responseStream /// private async ValueTask FormatOperationResultAsync( OperationResult operationResult, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); OperationResultFormatterContext? formatContext = null; var scope = Log.FormatOperationResultStart(); @@ -78,10 +79,11 @@ private async ValueTask FormatOperationResultAsync( /// private async ValueTask FormatResultBatchAsync( OperationResultBatch resultBatch, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); Exception? exception = null; using var semaphore = new SemaphoreSlim(1, 1); List? streams = null; @@ -120,6 +122,7 @@ private async ValueTask FormatResultBatchAsync( catch (Exception ex) { scope?.AddError(ex); + writer.Discard(); throw; } finally @@ -132,7 +135,7 @@ private async ValueTask FormatResultBatchAsync( } case IResponseStream responseStream: - keepAlive ??= new KeepAliveJob(semaphore, writer); + keepAlive ??= new KeepAliveJob(semaphore, pipe); var formatter = new StreamFormatter( _payloadFormatter, useIncrementalRfc1, @@ -176,15 +179,16 @@ private async ValueTask FormatResultBatchAsync( private async ValueTask FormatResponseStreamAsync( IResponseStream responseStream, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); using var semaphore = new SemaphoreSlim(1, 1); try { - using (var keepAlive = new KeepAliveJob(semaphore, writer)) + using (var keepAlive = new KeepAliveJob(semaphore, pipe)) { var formatter = new StreamFormatter( _payloadFormatter, @@ -234,7 +238,7 @@ private sealed class StreamFormatter( KeepAliveJob keepAliveJob, IResponseStream responseStream, SemaphoreSlim semaphore, - PipeWriter writer) + ChunkedPipeBufferWriter writer) { private OperationResultFormatterContext? _formatContext; @@ -263,6 +267,7 @@ public async Task ProcessAsync(CancellationToken ct) catch (Exception ex) { scope?.AddError(ex); + writer.Discard(); throw; } finally diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonResultFormatter.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonResultFormatter.cs index 58deb34f8c8..97d369d8a38 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonResultFormatter.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/JsonResultFormatter.cs @@ -236,19 +236,41 @@ private ValueTask FormatInternalAsync( ref OperationResultFormatterContext? context, CancellationToken cancellationToken) { - FormatInternal(result, writer, useIncrementalRfc1, ref context); - return FlushAsync(writer, cancellationToken); + var buffer = ChunkedPipeBufferWriter.Rent(writer); - static async ValueTask FlushAsync(PipeWriter w, CancellationToken ct) - => await w.FlushAsync(ct).ConfigureAwait(false); + try + { + FormatInternal(result, buffer, useIncrementalRfc1, ref context); + } + catch + { + buffer.Dispose(); + throw; + } + + return FlushAsync(buffer, cancellationToken); + + static async ValueTask FlushAsync(ChunkedPipeBufferWriter buffer, CancellationToken ct) + { + try + { + await buffer.FlushAsync(ct).ConfigureAwait(false); + } + finally + { + buffer.Dispose(); + } + } } private async ValueTask FormatInternalAsync( OperationResultBatch resultBatch, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken cancellationToken = default) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + foreach (var result in resultBatch.Results) { switch (result) @@ -300,10 +322,11 @@ private async ValueTask FormatInternalAsync( private async ValueTask FormatInternalAsync( IResponseStream responseStream, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken cancellationToken = default) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); OperationResultFormatterContext? context = null; try { diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/MultiPartResultFormatter.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/MultiPartResultFormatter.cs index 5c2a1147101..0ca62cf44b5 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Formatters/MultiPartResultFormatter.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/MultiPartResultFormatter.cs @@ -77,10 +77,11 @@ IResponseStream responseStream private async ValueTask FormatOperationResultAsync( OperationResult result, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct = default) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); OperationResultFormatterContext? formatContext = null; try { @@ -110,10 +111,11 @@ private async ValueTask FormatOperationResultAsync( private async ValueTask FormatResultBatchAsync( OperationResultBatch resultBatch, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct = default) { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); foreach (var result in resultBatch.Results) { switch (result) @@ -148,7 +150,7 @@ private async ValueTask FormatResultBatchAsync( break; case IResponseStream responseStream: - await FormatResponseStreamAsync( + await WriteResponseStreamAsync( responseStream, writer, useIncrementalRfc1, @@ -160,9 +162,19 @@ await FormatResponseStreamAsync( private async ValueTask FormatResponseStreamAsync( IResponseStream responseStream, - PipeWriter writer, + PipeWriter pipe, bool useIncrementalRfc1, CancellationToken ct = default) + { + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + await WriteResponseStreamAsync(responseStream, writer, useIncrementalRfc1, ct).ConfigureAwait(false); + } + + private async ValueTask WriteResponseStreamAsync( + IResponseStream responseStream, + ChunkedPipeBufferWriter writer, + bool useIncrementalRfc1, + CancellationToken ct) { // first, we create the iterator. await using var enumerator = responseStream.ReadResultsAsync().GetAsyncEnumerator(ct); diff --git a/src/HotChocolate/AspNetCore/src/Transport.Formatters/ThrowHelper.cs b/src/HotChocolate/AspNetCore/src/Transport.Formatters/ThrowHelper.cs new file mode 100644 index 00000000000..4ad8ebd3f36 --- /dev/null +++ b/src/HotChocolate/AspNetCore/src/Transport.Formatters/ThrowHelper.cs @@ -0,0 +1,13 @@ +namespace HotChocolate.Transport.Formatters; + +/// +/// An internal helper class that centralizes the exceptions of the formatters. +/// +internal static class ThrowHelper +{ + public static ArgumentOutOfRangeException ChunkedPipeBufferWriter_AdvanceNegative(int bytes) + => new(nameof(bytes), bytes, "The number of bytes to advance must not be negative."); + + public static InvalidOperationException ChunkedPipeBufferWriter_AdvancePastBuffer() + => new("Cannot advance past the end of the buffer."); +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ChunkedPipeBufferWriterTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ChunkedPipeBufferWriterTests.cs new file mode 100644 index 00000000000..985bce92dfb --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ChunkedPipeBufferWriterTests.cs @@ -0,0 +1,199 @@ +using System.IO.Pipelines; +using System.Text; +using HotChocolate.Transport.Formatters; + +namespace HotChocolate.AspNetCore.Serialization; + +public sealed class ChunkedPipeBufferWriterTests +{ + private const int ChunkSize = 64 * 1024; + + [Fact] + public async Task GetSpan_Should_BypassBuffer_When_SizeHintReachesDirectWriteThreshold() + { + // arrange + await using var output = new MemoryStream(); + var pipe = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + Fill(writer, 10, (byte)'a'); + + // act + const int size = ChunkedPipeBufferWriter.DirectWriteThreshold; + writer.GetSpan(size)[..size].Fill((byte)'b'); + writer.Advance(size); + await writer.FlushAsync(TestContext.Current.CancellationToken); + await pipe.CompleteAsync(); + + // assert + var expected = new byte[size + 10]; + expected.AsSpan(0, 10).Fill((byte)'a'); + expected.AsSpan(10).Fill((byte)'b'); + Assert.Equal(expected, output.ToArray()); + Assert.Equal(size, pipe.MaxSizeHint); + } + + [Fact] + public async Task GetSpan_Should_CopyBufferToPipe_When_SizeHintExceedsRemainingSpace() + { + // arrange + await using var output = new MemoryStream(); + var pipe = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + Fill(writer, ChunkSize - 10, (byte)'a'); + var spanCountBeforeOverflow = pipe.SpanCount; + + // act + writer.GetSpan(20)[..20].Fill((byte)'b'); + writer.Advance(20); + await writer.FlushAsync(TestContext.Current.CancellationToken); + await pipe.CompleteAsync(); + + // assert + var expected = new byte[ChunkSize + 10]; + expected.AsSpan(0, ChunkSize - 10).Fill((byte)'a'); + expected.AsSpan(ChunkSize - 10).Fill((byte)'b'); + Assert.Equal(expected, output.ToArray()); + Assert.Equal(0, spanCountBeforeOverflow); + Assert.Equal(0, pipe.MaxSizeHint); + } + + [Fact] + public void Advance_Should_Throw_When_BytesExceedBuffer() + { + // arrange + var pipe = PipeWriter.Create(Stream.Null); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + writer.GetSpan(5); + + // act + var exception = Assert.Throws(() => writer.Advance(ChunkSize + 1)); + + // assert + Assert.Equal("Cannot advance past the end of the buffer.", exception.Message); + } + + [Fact] + public void Advance_Should_Throw_When_BytesAreNegative() + { + // arrange + var pipe = PipeWriter.Create(Stream.Null); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + writer.GetSpan(5); + + // act + var exception = Assert.Throws(() => writer.Advance(-1)); + + // assert + Assert.Equal("bytes", exception.ParamName); + } + + [Fact] + public async Task FlushAsync_Should_ReleaseBuffer_When_Called() + { + // arrange + await using var output = new MemoryStream(); + var pipe = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + var hasBufferBeforeWrite = writer.HasBuffer; + Fill(writer, 5, (byte)'a'); + var hasBufferAfterWrite = writer.HasBuffer; + + // act + await writer.FlushAsync(TestContext.Current.CancellationToken); + await pipe.CompleteAsync(); + + // assert + Assert.False(hasBufferBeforeWrite); + Assert.True(hasBufferAfterWrite); + Assert.False(writer.HasBuffer); + Assert.Equal("aaaaa", Encoding.UTF8.GetString(output.ToArray())); + } + + [Fact] + public async Task Discard_Should_DropBufferedBytes_When_Called() + { + // arrange + await using var output = new MemoryStream(); + var pipe = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + Fill(writer, 3, (byte)'a'); + + // act + writer.Discard(); + var hasBufferAfterDiscard = writer.HasBuffer; + Fill(writer, 3, (byte)'b'); + await writer.FlushAsync(TestContext.Current.CancellationToken); + await pipe.CompleteAsync(); + + // assert + Assert.False(hasBufferAfterDiscard); + Assert.Equal("bbb", Encoding.UTF8.GetString(output.ToArray())); + } + + [Fact] + public async Task Discard_Should_KeepBytesCopiedToPipe_When_BufferOverflowed() + { + // arrange + await using var output = new MemoryStream(); + var pipe = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + using var writer = ChunkedPipeBufferWriter.Rent(pipe); + Fill(writer, ChunkSize - 10, (byte)'a'); + // the overflowing write copies the buffered bytes to the pipe + Fill(writer, 20, (byte)'b'); + + // act + writer.Discard(); + await writer.FlushAsync(TestContext.Current.CancellationToken); + await pipe.CompleteAsync(); + + // assert + var expected = new byte[ChunkSize - 10]; + expected.AsSpan().Fill((byte)'a'); + Assert.Equal(expected, output.ToArray()); + } + + [Fact] + public async Task Dispose_Should_DiscardUnflushedBytes_When_Called() + { + // arrange + await using var output = new MemoryStream(); + var pipe = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + var writer = ChunkedPipeBufferWriter.Rent(pipe); + Fill(writer, 5, (byte)'a'); + + // act + writer.Dispose(); + await pipe.CompleteAsync(); + + // assert + Assert.Empty(output.ToArray()); + } + + [Fact] + public void Rent_Should_ReuseDisposedWriter_When_RentedOnSameThread() + { + // arrange + var pipe = PipeWriter.Create(Stream.Null); + var first = ChunkedPipeBufferWriter.Rent(pipe); + first.Dispose(); + + // act + using var second = ChunkedPipeBufferWriter.Rent(pipe); + + // assert + Assert.Same(first, second); + } + + /// + /// Writes bytes of in writes of at most 1 KB. + /// + private static void Fill(ChunkedPipeBufferWriter writer, int count, byte value) + { + for (var written = 0; written < count; written += 1024) + { + var size = Math.Min(1024, count - written); + writer.GetSpan(size)[..size].Fill(value); + writer.Advance(size); + } + } +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/CountingPipeWriter.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/CountingPipeWriter.cs new file mode 100644 index 00000000000..390dfa135ab --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/CountingPipeWriter.cs @@ -0,0 +1,59 @@ +using System.IO.Pipelines; + +namespace HotChocolate.AspNetCore.Serialization; + +/// +/// A pipe writer that forwards to and counts the calls it receives. +/// +internal sealed class CountingPipeWriter(PipeWriter inner) : PipeWriter +{ + /// + /// Gets the number of and calls. + /// + public int SpanCount { get; private set; } + + /// + /// Gets the largest size hint passed to or . + /// + public int MaxSizeHint { get; private set; } + + /// + /// Gets the number of calls. + /// + public int AdvanceCount { get; private set; } + + /// + /// Gets the number of calls. + /// + public int FlushCount { get; private set; } + + public override void Advance(int bytes) + { + AdvanceCount++; + inner.Advance(bytes); + } + + public override Memory GetMemory(int sizeHint = 0) + { + SpanCount++; + MaxSizeHint = Math.Max(MaxSizeHint, sizeHint); + return inner.GetMemory(sizeHint); + } + + public override Span GetSpan(int sizeHint = 0) + { + SpanCount++; + MaxSizeHint = Math.Max(MaxSizeHint, sizeHint); + return inner.GetSpan(sizeHint); + } + + public override ValueTask FlushAsync(CancellationToken cancellationToken = default) + { + FlushCount++; + return inner.FlushAsync(cancellationToken); + } + + public override void CancelPendingFlush() => inner.CancelPendingFlush(); + + public override void Complete(Exception? exception = null) => inner.Complete(exception); +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/EventStreamResultFormatterTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/EventStreamResultFormatterTests.cs new file mode 100644 index 00000000000..beaff7a432b --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/EventStreamResultFormatterTests.cs @@ -0,0 +1,81 @@ +using System.IO.Pipelines; +using System.Text; +using HotChocolate.Execution; +using HotChocolate.Transport.Formatters; +using static HotChocolate.AspNetCore.Serialization.ResultTestHelper; + +namespace HotChocolate.AspNetCore.Serialization; + +public sealed class EventStreamResultFormatterTests +{ + [Fact] + public async Task FormatAsync_Should_DiscardPartialMessage_When_BatchResultFormattingThrows() + { + // arrange + await using var output = new MemoryStream(); + var writer = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + var batch = new OperationResultBatch( + [ + CreateResult(new Dictionary { ["a"] = 1 }), + CreateResult(jsonWriter => + { + jsonWriter.WriteStartObject(); + jsonWriter.WritePropertyName("partial"); + throw new InvalidOperationException("boom"); + }), + CreateResult(new Dictionary { ["c"] = 3 }) + ]); + + // act + await new EventStreamResultFormatter(default).FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + TestContext.Current.CancellationToken); + await writer.CompleteAsync(); + + // assert + Assert.Equal( + "event: next\ndata: {\"data\":{\"a\":1}}\n\n" + + "event: next\ndata: {\"data\":{\"c\":3}}\n\n" + + "event: complete\ndata:\n\n", + Encoding.UTF8.GetString(output.ToArray())); + } + + [Fact] + public async Task FormatAsync_Should_DiscardPartialMessage_When_StreamResultFormattingThrows() + { + // arrange + await using var output = new MemoryStream(); + var writer = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + // the healthy stream yields only after the failing stream wrote its partial message + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var batch = new OperationResultBatch( + [ + CreateStream( + CreateResult(jsonWriter => + { + jsonWriter.WriteStartObject(); + jsonWriter.WritePropertyName("partial"); + gate.SetResult(); + throw new InvalidOperationException("boom"); + })), + CreateStream(gate.Task, CreateResult(new Dictionary { ["s2"] = 1 })) + ]); + + // act + var exception = await Assert.ThrowsAsync( + async () => await new EventStreamResultFormatter(default).FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + TestContext.Current.CancellationToken)); + await writer.CompleteAsync(); + + // assert + Assert.Equal("boom", exception.Message); + Assert.Equal( + "event: next\ndata: {\"data\":{\"s2\":1}}\n\n", + Encoding.UTF8.GetString(output.ToArray())); + } +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonLinesResultFormatterTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonLinesResultFormatterTests.cs new file mode 100644 index 00000000000..8a5ec599168 --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonLinesResultFormatterTests.cs @@ -0,0 +1,79 @@ +using System.IO.Pipelines; +using System.Text; +using HotChocolate.Execution; +using HotChocolate.Transport.Formatters; +using static HotChocolate.AspNetCore.Serialization.ResultTestHelper; + +namespace HotChocolate.AspNetCore.Serialization; + +public sealed class JsonLinesResultFormatterTests +{ + [Fact] + public async Task FormatAsync_Should_DiscardPartialResult_When_BatchResultFormattingThrows() + { + // arrange + await using var output = new MemoryStream(); + var writer = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + // the stream yields only after the failing result wrote its partial bytes + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var batch = new OperationResultBatch( + [ + CreateStream(gate.Task, CreateResult(new Dictionary { ["s"] = 1 })), + CreateResult(jsonWriter => + { + jsonWriter.WriteStartObject(); + jsonWriter.WritePropertyName("partial"); + gate.SetResult(); + throw new InvalidOperationException("boom"); + }) + ]); + + // act + var exception = await Assert.ThrowsAsync( + async () => await new JsonLinesResultFormatter(default).FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + TestContext.Current.CancellationToken)); + await writer.CompleteAsync(); + + // assert + Assert.Equal("boom", exception.Message); + Assert.Equal("{\"data\":{\"s\":1}}\n", Encoding.UTF8.GetString(output.ToArray())); + } + + [Fact] + public async Task FormatAsync_Should_DiscardPartialResult_When_StreamResultFormattingThrows() + { + // arrange + await using var output = new MemoryStream(); + var writer = PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true)); + // the healthy stream yields only after the failing stream wrote its partial bytes + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var batch = new OperationResultBatch( + [ + CreateStream( + CreateResult(jsonWriter => + { + jsonWriter.WriteStartObject(); + jsonWriter.WritePropertyName("partial"); + gate.SetResult(); + throw new InvalidOperationException("boom"); + })), + CreateStream(gate.Task, CreateResult(new Dictionary { ["s2"] = 1 })) + ]); + + // act + var exception = await Assert.ThrowsAsync( + async () => await new JsonLinesResultFormatter(default).FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + TestContext.Current.CancellationToken)); + await writer.CompleteAsync(); + + // assert + Assert.Equal("boom", exception.Message); + Assert.Equal("{\"data\":{\"s2\":1}}\n", Encoding.UTF8.GetString(output.ToArray())); + } +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonResultFormatterTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonResultFormatterTests.cs new file mode 100644 index 00000000000..456d59c7e47 --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/JsonResultFormatterTests.cs @@ -0,0 +1,157 @@ +using System.Buffers; +using System.IO.Pipelines; +using System.Text; +using HotChocolate.Execution; +using HotChocolate.Transport.Formatters; +using static HotChocolate.AspNetCore.Serialization.ResultTestHelper; + +namespace HotChocolate.AspNetCore.Serialization; + +public sealed class JsonResultFormatterTests +{ + [Fact] + public async Task FormatAsync_Should_WriteResultInOneChunk_When_ResultFitsIntoChunk() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var result = new OperationResult( + CreateData( + new Dictionary + { + ["product"] = new Dictionary + { + ["name"] = "Abc", + ["price"] = 12.5 + } + })); + + // act + await JsonResultFormatter.Default.FormatAsync( + result, + writer, + ExecutionResultFormatFlags.None, + CancellationToken.None); + await writer.CompleteAsync(); + + // assert + Assert.Equal( + """{"data":{"product":{"name":"Abc","price":12.5}}}""", + Encoding.UTF8.GetString(output.ToArray())); + Assert.Equal(1, writer.SpanCount); + Assert.Equal(1, writer.AdvanceCount); + Assert.Equal(1, writer.FlushCount); + } + + [Fact] + public async Task FormatAsync_Should_WriteCompleteResult_When_ResultExceedsChunkSize() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var items = new List(); + + for (var i = 0; i < 4000; i++) + { + items.Add(new Dictionary { ["id"] = i, ["name"] = $"item-{i}-{new string('x', 40)}" }); + } + + var data = new Dictionary { ["items"] = items }; + var result = new OperationResult(CreateData(data)); + var expected = new ArrayBufferWriter(); + JsonResultFormatter.Default.Format(new OperationResult(CreateData(data)), expected); + + // act + await JsonResultFormatter.Default.FormatAsync( + result, + writer, + ExecutionResultFormatFlags.None, + CancellationToken.None); + await writer.CompleteAsync(); + + // assert + Assert.True(expected.WrittenCount > 64 * 1024); + Assert.Equal(expected.WrittenSpan.ToArray(), output.ToArray()); + Assert.Equal(1, writer.FlushCount); + } + + [Fact] + public async Task FormatAsync_Should_WriteScalarThroughPipe_When_ScalarExceedsChunkSize() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var text = new string('x', 1024 * 1024); + var data = new Dictionary { ["text"] = text, ["after"] = "tail" }; + var result = new OperationResult(CreateData(data)); + var expected = new ArrayBufferWriter(); + JsonResultFormatter.Default.Format(new OperationResult(CreateData(data)), expected); + + // act + await JsonResultFormatter.Default.FormatAsync( + result, + writer, + ExecutionResultFormatFlags.None, + CancellationToken.None); + await writer.CompleteAsync(); + + // assert + Assert.Equal(expected.WrittenSpan.ToArray(), output.ToArray()); + Assert.Equal(text.Length * 3 + 3, writer.MaxSizeHint); + Assert.Equal(1, writer.FlushCount); + } + + [Fact] + public async Task FormatAsync_Should_WriteBatchWithOneFlush_When_ResultIsBatch() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var batch = new OperationResultBatch( + [ + new OperationResult(CreateData(new Dictionary { ["a"] = 1 })), + new OperationResult(CreateData(new Dictionary { ["b"] = 2 })) + ]); + + // act + await JsonResultFormatter.Default.FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + CancellationToken.None); + await writer.CompleteAsync(); + + // assert + Assert.Equal("""{"data":{"a":1}}{"data":{"b":2}}""", Encoding.UTF8.GetString(output.ToArray())); + Assert.Equal(1, writer.SpanCount); + Assert.Equal(1, writer.FlushCount); + } + + [Fact] + public async Task FormatAsync_Should_FlushPerResult_When_ResultIsStream() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var stream = new ResponseStream( + () => CreateResults( + [ + new OperationResult(CreateData(new Dictionary { ["a"] = 1 })), + new OperationResult(CreateData(new Dictionary { ["b"] = 2 })) + ]), + ExecutionResultKind.SubscriptionResult); + + // act + await JsonResultFormatter.Default.FormatAsync( + stream, + writer, + ExecutionResultFormatFlags.None, + CancellationToken.None); + await writer.CompleteAsync(); + + // assert + Assert.Equal("""{"data":{"a":1}}{"data":{"b":2}}""", Encoding.UTF8.GetString(output.ToArray())); + Assert.Equal(2, writer.SpanCount); + Assert.Equal(2, writer.FlushCount); + } +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/MultiPartResultFormatterTests.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/MultiPartResultFormatterTests.cs new file mode 100644 index 00000000000..a1a1609a7aa --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/MultiPartResultFormatterTests.cs @@ -0,0 +1,46 @@ +using System.IO.Pipelines; +using System.Text; +using HotChocolate.Execution; +using HotChocolate.Transport.Formatters; +using static HotChocolate.AspNetCore.Serialization.ResultTestHelper; + +namespace HotChocolate.AspNetCore.Serialization; + +public sealed class MultiPartResultFormatterTests +{ + private const string Header = "Content-Type: application/json; charset=utf-8\r\n\r\n"; + + [Fact] + public async Task FormatAsync_Should_WriteStreamBetweenResults_When_BatchContainsStream() + { + // arrange + await using var output = new MemoryStream(); + var writer = new CountingPipeWriter(PipeWriter.Create(output, new StreamPipeWriterOptions(leaveOpen: true))); + var batch = new OperationResultBatch( + [ + CreateResult(new Dictionary { ["a"] = 1 }), + CreateStream( + CreateResult(new Dictionary { ["b"] = 2 }), + CreateResult(new Dictionary { ["c"] = 3 })), + CreateResult(new Dictionary { ["d"] = 4 }) + ]); + + // act + await new MultiPartResultFormatter().FormatAsync( + batch, + writer, + ExecutionResultFormatFlags.None, + TestContext.Current.CancellationToken); + await writer.CompleteAsync(); + + // assert + const string expected = + "\r\n---\r\n" + Header + """{"data":{"a":1}}""" + "\r\n-----\r\n" + + "\r\n---\r\n" + Header + """{"data":{"b":2}}""" + + "\r\n---\r\n" + Header + """{"data":{"c":3}}""" + + "\r\n-----\r\n" + + "\r\n---\r\n" + Header + """{"data":{"d":4}}""" + "\r\n-----\r\n"; + Assert.Equal(expected, Encoding.UTF8.GetString(output.ToArray())); + Assert.Equal(5, writer.FlushCount); + } +} diff --git a/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ResultTestHelper.cs b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ResultTestHelper.cs new file mode 100644 index 00000000000..df52050394d --- /dev/null +++ b/src/HotChocolate/AspNetCore/test/AspNetCore.Tests/Serialization/ResultTestHelper.cs @@ -0,0 +1,66 @@ +using System.Text.Json; +using HotChocolate.Execution; +using HotChocolate.Text.Json; + +namespace HotChocolate.AspNetCore.Serialization; + +/// +/// Creates results, batches and streams for formatter tests. +/// +internal static class ResultTestHelper +{ + private static readonly JsonSerializerOptions s_options = new(JsonSerializerDefaults.Web); + + public static OperationResultData CreateData(object value) + => new(value, isValueNull: false, new DictionaryJsonFormatter(value), memoryHolder: null); + + public static OperationResult CreateResult(Dictionary data) + => new(CreateData(data)); + + /// + /// Creates a result whose data is written by . An exception thrown + /// by escapes from the formatter. + /// + public static OperationResult CreateResult(Action write) + => new(new OperationResultData(write, isValueNull: false, new DelegateJsonFormatter(write), memoryHolder: null)); + + public static ResponseStream CreateStream(params OperationResult[] results) + => new(() => CreateResults(results), ExecutionResultKind.SubscriptionResult); + + /// + /// Creates a stream that yields its first result only after completed. + /// + public static ResponseStream CreateStream(Task gate, params OperationResult[] results) + => new(() => CreateResults(gate, results), ExecutionResultKind.SubscriptionResult); + + public static async IAsyncEnumerable CreateResults(IReadOnlyList results) + { + foreach (var result in results) + { + yield return result; + await Task.Yield(); + } + } + + private static async IAsyncEnumerable CreateResults(Task gate, IReadOnlyList results) + { + await gate; + + foreach (var result in results) + { + yield return result; + await Task.Yield(); + } + } + + private sealed class DictionaryJsonFormatter(object value) : IRawJsonFormatter + { + public void WriteDataTo(JsonWriter jsonWriter) + => JsonValueFormatter.WriteValue(jsonWriter, value, s_options); + } + + private sealed class DelegateJsonFormatter(Action write) : IRawJsonFormatter + { + public void WriteDataTo(JsonWriter jsonWriter) => write(jsonWriter); + } +}