From 1ab707796fc2ebbb6f592e8f0fd1038e3f0ed4f6 Mon Sep 17 00:00:00 2001 From: Sam Xu Date: Sun, 23 Aug 2026 14:35:24 -0700 Subject: [PATCH] Prevent duplicate pooled buffer returns Remove duplicate stream ownership in ODataJsonValueSerializer. EndStreamValueScope and EndStreamValueScopeAsync are now solely responsible for disposing writer-owned streams, preventing the serializer and writer from returning the same pooled byte array twice. Make binary stream and text writer disposal idempotent across synchronous, asynchronous, repeated, and mixed disposal paths. Interlocked disposal gates and atomic ownership detachment guarantee that each rented byte or character array is returned at most once and that repeated scope completion cannot reuse an already-disposed writer. This prevents duplicate references from entering a shared ArrayPool, where simultaneous later rentals could otherwise alias the same array and cause cross-operation data corruption or disclosure. Use nested try/finally cleanup so writer scopes are completed when copying fails and source streams still honor ODataBinaryStreamValue.LeaveOpen. Preserve ConfigureAwait(false) for asynchronous library cleanup. Add internal injectable array pools for deterministic ownership tests while production constructors continue using the shared pools. Cover exact rent/return accounting, repeated and mixed disposal, serializer ownership, and Base64 boundary payload sizes 1, 2, 3, 2048, and 2049. --- .../Json/ODataJsonValueSerializer.cs | 50 ++++++---- .../Json/ODataUtf8JsonWriter.Stream.cs | 96 +++++++++++++------ .../Json/ODataUtf8JsonWriter.TextWriter.cs | 93 +++++++++++++----- .../Json/MockJsonWriter.cs | 12 ++- .../ODataJsonValueSerializerAsyncTests.cs | 35 +++++++ .../Json/ODataJsonValueSerializerTests.cs | 21 ++++ .../Json/ODataUtf8JsonTextWriterTests.cs | 20 ++++ .../Json/ODataUtf8JsonWriterStreamTests.cs | 71 ++++++++++++++ 8 files changed, 324 insertions(+), 74 deletions(-) diff --git a/src/Microsoft.OData.Core/Json/ODataJsonValueSerializer.cs b/src/Microsoft.OData.Core/Json/ODataJsonValueSerializer.cs index b17281abe0..9c78cb3036 100644 --- a/src/Microsoft.OData.Core/Json/ODataJsonValueSerializer.cs +++ b/src/Microsoft.OData.Core/Json/ODataJsonValueSerializer.cs @@ -382,15 +382,24 @@ public virtual void WriteUntypedValue( public virtual void WriteStreamValue(ODataBinaryStreamValue streamValue) { - Stream stream = this.JsonWriter.StartStreamValueScope(); - streamValue.Stream.CopyTo(stream); - stream.Flush(); - stream.Dispose(); - this.JsonWriter.EndStreamValueScope(); - - if (!streamValue.LeaveOpen) + try { - streamValue.Stream.Dispose(); + Stream stream = this.JsonWriter.StartStreamValueScope(); + try + { + streamValue.Stream.CopyTo(stream); + } + finally + { + this.JsonWriter.EndStreamValueScope(); + } + } + finally + { + if (!streamValue.LeaveOpen) + { + streamValue.Stream.Dispose(); + } } } @@ -722,15 +731,24 @@ public virtual Task WriteUntypedValueAsync( /// A task that represents the asynchronous write operation. public virtual async Task WriteStreamValueAsync(ODataBinaryStreamValue streamValue) { - Stream stream = await this.JsonWriter.StartStreamValueScopeAsync().ConfigureAwait(false); - await streamValue.Stream.CopyToAsync(stream).ConfigureAwait(false); - await stream.FlushAsync().ConfigureAwait(false); - await stream.DisposeAsync().ConfigureAwait(false); - await this.JsonWriter.EndStreamValueScopeAsync().ConfigureAwait(false); - - if (!streamValue.LeaveOpen) + try { - await streamValue.Stream.DisposeAsync().ConfigureAwait(false); + Stream stream = await this.JsonWriter.StartStreamValueScopeAsync().ConfigureAwait(false); + try + { + await streamValue.Stream.CopyToAsync(stream).ConfigureAwait(false); + } + finally + { + await this.JsonWriter.EndStreamValueScopeAsync().ConfigureAwait(false); + } + } + finally + { + if (!streamValue.LeaveOpen) + { + await streamValue.Stream.DisposeAsync().ConfigureAwait(false); + } } } diff --git a/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.Stream.cs b/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.Stream.cs index e893d82fe3..86fdf55d27 100644 --- a/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.Stream.cs +++ b/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.Stream.cs @@ -43,8 +43,9 @@ public Stream StartStreamValueScope() /// public void EndStreamValueScope() { - this.binaryValueStream?.Dispose(); + Stream stream = this.binaryValueStream; this.binaryValueStream = null; + stream?.Dispose(); this.Flush(); this.bufferWriter.Write(this.DoubleQuote.Slice(0, 1).Span); @@ -76,10 +77,11 @@ public async Task StartStreamValueScopeAsync() /// A task representing the asynchronous operation. public async Task EndStreamValueScopeAsync() { - if (this.binaryValueStream != null) + Stream stream = this.binaryValueStream; + this.binaryValueStream = null; + if (stream != null) { - await this.binaryValueStream.DisposeAsync().ConfigureAwait(false); - this.binaryValueStream = null; + await stream.DisposeAsync().ConfigureAwait(false); } await this.DrainBufferIfThresholdReachedAsync().ConfigureAwait(false); @@ -96,16 +98,24 @@ public async Task EndStreamValueScopeAsync() internal sealed class ODataUtf8JsonWriteStream : Stream { private readonly ODataUtf8JsonWriter jsonWriter = null; + private readonly ArrayPool arrayPool; private byte[] buffer; - int numBytesNotWrittenFromPreviousChunk = 0; + private int disposed; + private int numBytesNotWrittenFromPreviousChunk = 0; /// /// Initializes a new instance of the class with the specified ODataUtf8JsonWriter. /// /// The OData UTF-8 JSON writer to write to. internal ODataUtf8JsonWriteStream(ODataUtf8JsonWriter writer) + : this(writer, ArrayPool.Shared) + { + } + + internal ODataUtf8JsonWriteStream(ODataUtf8JsonWriter writer, ArrayPool arrayPool) { this.jsonWriter = writer; + this.arrayPool = arrayPool; } public override bool CanRead => false; @@ -147,22 +157,35 @@ public override async Task FlushAsync(CancellationToken cancellationToken) /// true if called from Dispose; false if called form the finalizer. protected override void Dispose(bool disposing) { - if (this.numBytesNotWrittenFromPreviousChunk > 0) + if (Interlocked.Exchange(ref this.disposed, 1) != 0) { - // If there are unprocessed bytes, encode and write them as the final block. - ReadOnlySpan bytesNotProcessedFromPreviousChunk = this.buffer.AsSpan().Slice(0, this.numBytesNotWrittenFromPreviousChunk); - - this.jsonWriter.Base64EncodeAndWriteChunk(bytesNotProcessedFromPreviousChunk, isFinalBlock: true, out this.numBytesNotWrittenFromPreviousChunk); - Debug.Assert(numBytesNotWrittenFromPreviousChunk == 0, "numBytesNotWrittenFromPreviousChunk == 0"); + return; } - if (this.buffer != null) + byte[] buffer = Interlocked.Exchange(ref this.buffer, null); + + try { - ArrayPool.Shared.Return(this.buffer); + if (buffer != null && this.numBytesNotWrittenFromPreviousChunk > 0) + { + // If there are unprocessed bytes, encode and write them as the final block. + ReadOnlySpan bytesNotProcessedFromPreviousChunk = buffer.AsSpan().Slice(0, this.numBytesNotWrittenFromPreviousChunk); + + this.jsonWriter.Base64EncodeAndWriteChunk(bytesNotProcessedFromPreviousChunk, isFinalBlock: true, out this.numBytesNotWrittenFromPreviousChunk); + Debug.Assert(numBytesNotWrittenFromPreviousChunk == 0, "numBytesNotWrittenFromPreviousChunk == 0"); + } + + this.Flush(); } + finally + { + if (buffer != null) + { + this.arrayPool.Return(buffer); + } - this.Flush(); - base.Dispose(disposing); + base.Dispose(disposing); + } } /// @@ -171,21 +194,36 @@ protected override void Dispose(bool disposing) /// A representing the asynchronous operation. public override async ValueTask DisposeAsync() { - if (this.numBytesNotWrittenFromPreviousChunk > 0) + if (Interlocked.Exchange(ref this.disposed, 1) != 0) { - // If there are unprocessed bytes, encode and write them as the final block. - ReadOnlyMemory bytesNotProcessedFromPreviousChunk = this.buffer.AsMemory().Slice(0, this.numBytesNotWrittenFromPreviousChunk); - - this.jsonWriter.Base64EncodeAndWriteChunk(bytesNotProcessedFromPreviousChunk.Span, isFinalBlock: true, out this.numBytesNotWrittenFromPreviousChunk); - Debug.Assert(numBytesNotWrittenFromPreviousChunk == 0, "numBytesNotWrittenFromPreviousChunk == 0"); + return; } - if (this.buffer != null) + byte[] buffer = Interlocked.Exchange(ref this.buffer, null); + + try { - ArrayPool.Shared.Return(this.buffer); + if (buffer != null && this.numBytesNotWrittenFromPreviousChunk > 0) + { + // If there are unprocessed bytes, encode and write them as the final block. + ReadOnlyMemory bytesNotProcessedFromPreviousChunk = buffer.AsMemory().Slice(0, this.numBytesNotWrittenFromPreviousChunk); + + this.jsonWriter.Base64EncodeAndWriteChunk(bytesNotProcessedFromPreviousChunk.Span, isFinalBlock: true, out this.numBytesNotWrittenFromPreviousChunk); + Debug.Assert(numBytesNotWrittenFromPreviousChunk == 0, "numBytesNotWrittenFromPreviousChunk == 0"); + } + + await this.jsonWriter.FlushAsync().ConfigureAwait(false); } + finally + { + if (buffer != null) + { + this.arrayPool.Return(buffer); + } - await this.jsonWriter.FlushAsync().ConfigureAwait(false); + base.Dispose(true); + GC.SuppressFinalize(this); + } } /// @@ -263,7 +301,7 @@ private void WriteByteValueInChunks(ReadOnlySpan value) int totalLength = bytesNotProcessedFromPreviousChunk.Length + chunk.Length; // Rent an array to hold bytes from both previous and current chunks. - byte[] combinedArray = ArrayPool.Shared.Rent(totalLength); + byte[] combinedArray = this.arrayPool.Rent(totalLength); // Copy bytes from bytesNotProcessedFromPreviousChunk to the combined array bytesNotProcessedFromPreviousChunk.CopyTo(combinedArray); @@ -273,7 +311,7 @@ private void WriteByteValueInChunks(ReadOnlySpan value) WriteChunk(combinedArray.AsSpan().Slice(0, totalLength), totalLength, isFinalBlock); - ArrayPool.Shared.Return(combinedArray); + this.arrayPool.Return(combinedArray); } else { @@ -309,7 +347,7 @@ private async ValueTask WriteByteValueInChunksAsync(ReadOnlyMemory value) int totalLength = bytesNotProcessedFromPreviousChunk.Length + chunk.Length; // Rent an array to hold bytes from both previous and current chunks. - byte[] combinedArray = ArrayPool.Shared.Rent(totalLength); + byte[] combinedArray = this.arrayPool.Rent(totalLength); // Copy bytes from bytesNotProcessedFromPreviousChunk to the combined array bytesNotProcessedFromPreviousChunk.Span.CopyTo(combinedArray); @@ -319,7 +357,7 @@ private async ValueTask WriteByteValueInChunksAsync(ReadOnlyMemory value) WriteChunk(combinedArray.AsSpan().Slice(0, totalLength), totalLength, isFinalBlock); - ArrayPool.Shared.Return(combinedArray); + this.arrayPool.Return(combinedArray); } else { @@ -345,7 +383,7 @@ private void WriteChunk(ReadOnlySpan chunk, int chunkLength, bool isFinalB { if (this.buffer == null) { - this.buffer = ArrayPool.Shared.Rent(ODataUtf8JsonWriter.chunkSize); + this.buffer = this.arrayPool.Rent(ODataUtf8JsonWriter.chunkSize); } // Update the buffer with unprocessed bytes from the current chunk. diff --git a/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.TextWriter.cs b/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.TextWriter.cs index b6d4140eca..36204f42b2 100644 --- a/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.TextWriter.cs +++ b/src/Microsoft.OData.Core/Json/ODataUtf8JsonWriter.TextWriter.cs @@ -11,6 +11,7 @@ namespace Microsoft.OData.Json using System.IO; using System.Text; using System.Text.Unicode; + using System.Threading; using System.Threading.Tasks; internal sealed partial class ODataUtf8JsonWriter @@ -68,7 +69,9 @@ public void EndTextWriterValueScope() this.bufferWriter.Write(this.DoubleQuote.Slice(0, 1).Span); } - this.textWriter?.Dispose(); + TextWriter textWriter = this.textWriter; + this.textWriter = null; + textWriter?.Dispose(); this.Flush(); CheckIfSeparatorNeeded(); @@ -113,9 +116,11 @@ public async Task EndTextWriterValueScopeAsync() this.bufferWriter.Write(this.DoubleQuote.Slice(0, 1).Span); } - if (this.textWriter != null) + TextWriter textWriter = this.textWriter; + this.textWriter = null; + if (textWriter != null) { - await this.textWriter.DisposeAsync().ConfigureAwait(false); + await textWriter.DisposeAsync().ConfigureAwait(false); } await this.DrainBufferIfThresholdReachedAsync().ConfigureAwait(false); @@ -139,6 +144,7 @@ private static bool CheckIfWritingJson(string currentContentType) internal sealed class ODataUtf8JsonTextWriter : TextWriter { private readonly ODataUtf8JsonWriter jsonWriter = null; + private readonly ArrayPool arrayPool; // Buffer used to store chars that could not be encoded due to // insufficient data in the input buffer. The chars will be prepended // to the next chunk of input. @@ -148,10 +154,17 @@ internal sealed class ODataUtf8JsonTextWriter : TextWriter // This buffer is used by Write(char) to store the char so // that we can re-use our Write(char[], ...) method. private char[] singleCharBuffer; + private int disposed; public ODataUtf8JsonTextWriter(ODataUtf8JsonWriter jsonWriter) + : this(jsonWriter, ArrayPool.Shared) + { + } + + internal ODataUtf8JsonTextWriter(ODataUtf8JsonWriter jsonWriter, ArrayPool arrayPool) { this.jsonWriter = jsonWriter; + this.arrayPool = arrayPool; } /// @@ -182,18 +195,32 @@ public override async Task FlushAsync() /// true if called from Dispose; false if called form the finalizer. protected override void Dispose(bool disposing) { - if (this.buffer != null) + if (Interlocked.Exchange(ref this.disposed, 1) != 0) { - ArrayPool.Shared.Return(this.buffer); + return; } - if (this.singleCharBuffer != null) + char[] buffer = Interlocked.Exchange(ref this.buffer, null); + char[] singleCharBuffer = Interlocked.Exchange(ref this.singleCharBuffer, null); + + try { - ArrayPool.Shared.Return(this.singleCharBuffer); + this.Flush(); } + finally + { + if (buffer != null) + { + this.arrayPool.Return(buffer); + } + + if (singleCharBuffer != null) + { + this.arrayPool.Return(singleCharBuffer); + } - this.Flush(); - base.Dispose(disposing); + base.Dispose(disposing); + } } /// @@ -202,17 +229,33 @@ protected override void Dispose(bool disposing) /// A representing the asynchronous operation. public override async ValueTask DisposeAsync() { - if (this.buffer != null) + if (Interlocked.Exchange(ref this.disposed, 1) != 0) { - ArrayPool.Shared.Return(this.buffer); + return; } - if (this.singleCharBuffer != null) + char[] buffer = Interlocked.Exchange(ref this.buffer, null); + char[] singleCharBuffer = Interlocked.Exchange(ref this.singleCharBuffer, null); + + try { - ArrayPool.Shared.Return(this.singleCharBuffer); + await this.FlushAsync().ConfigureAwait(false); } + finally + { + if (buffer != null) + { + this.arrayPool.Return(buffer); + } - await this.FlushAsync().ConfigureAwait(false); + if (singleCharBuffer != null) + { + this.arrayPool.Return(singleCharBuffer); + } + + base.Dispose(true); + GC.SuppressFinalize(this); + } } /// @@ -238,7 +281,7 @@ public override void Write(char value) /// The character to write. public override async Task WriteAsync(char value) { - this.singleCharBuffer ??= ArrayPool.Shared.Rent(1); + this.singleCharBuffer ??= this.arrayPool.Rent(1); this.singleCharBuffer[0] = value; ReadOnlyMemory input = this.singleCharBuffer.AsMemory().Slice(0, 1); @@ -314,7 +357,7 @@ private void WriteCharsInChunks(ReadOnlySpan value) ReadOnlySpan charsNotProcessedFromPreviousChunk = this.buffer.AsSpan().Slice(0, this.numOfCharsNotWrittenFromPreviousChunk); int totalLength = charsNotProcessedFromPreviousChunk.Length + chunk.Length; - char[] combinedArray = ArrayPool.Shared.Rent(totalLength); + char[] combinedArray = this.arrayPool.Rent(totalLength); // Copy chars from charsNotProcessedFromPreviousChunk to the combined array charsNotProcessedFromPreviousChunk.CopyTo(combinedArray); @@ -327,7 +370,7 @@ private void WriteCharsInChunks(ReadOnlySpan value) // Write the chunk. this.WriteChunk(combinedArray.AsSpan().Slice(0, totalLength), totalLength, firstIndexToEscape, isFinalBlock); - ArrayPool.Shared.Return(combinedArray); + this.arrayPool.Return(combinedArray); } else { @@ -362,7 +405,7 @@ private async ValueTask WriteCharsInChunksAsync(ReadOnlyMemory value) ReadOnlyMemory charsNotProcessedFromPreviousChunk = this.buffer.AsMemory().Slice(0, this.numOfCharsNotWrittenFromPreviousChunk); int totalLength = charsNotProcessedFromPreviousChunk.Length + chunk.Length; - char[] combinedArray = ArrayPool.Shared.Rent(totalLength); + char[] combinedArray = this.arrayPool.Rent(totalLength); // Copy chars from charsNotProcessedFromPreviousChunk to the combined array charsNotProcessedFromPreviousChunk.CopyTo(combinedArray); @@ -375,7 +418,7 @@ private async ValueTask WriteCharsInChunksAsync(ReadOnlyMemory value) // Write the chunk. this.WriteChunk(combinedArray.AsSpan().Slice(0, totalLength), totalLength, firstIndexToEscape, isFinalBlock); - ArrayPool.Shared.Return(combinedArray); + this.arrayPool.Return(combinedArray); } else { @@ -411,7 +454,7 @@ private void WriteCharsInChunksWithoutEscaping(ReadOnlySpan value) ReadOnlySpan charsNotProcessedFromPreviousChunk = this.buffer.AsSpan().Slice(0, this.numOfCharsNotWrittenFromPreviousChunk); int totalLength = charsNotProcessedFromPreviousChunk.Length + chunk.Length; - char[] combinedArray = ArrayPool.Shared.Rent(totalLength); + char[] combinedArray = this.arrayPool.Rent(totalLength); // Copy chars from charsNotProcessedFromPreviousChunk to the combined array charsNotProcessedFromPreviousChunk.CopyTo(combinedArray); @@ -422,7 +465,7 @@ private void WriteCharsInChunksWithoutEscaping(ReadOnlySpan value) // Write the chunk. this.WriteChunkWithoutEscaping(combinedArray.AsSpan().Slice(0, totalLength), totalLength, isFinalBlock); - ArrayPool.Shared.Return(combinedArray, true); + this.arrayPool.Return(combinedArray, true); } else { @@ -455,7 +498,7 @@ private async ValueTask WriteCharsInChunksWithoutEscapingAsync(ReadOnlyMemory charsNotProcessedFromPreviousChunk = this.buffer.AsMemory().Slice(0, this.numOfCharsNotWrittenFromPreviousChunk); int totalLength = charsNotProcessedFromPreviousChunk.Length + chunk.Length; - char[] combinedArray = ArrayPool.Shared.Rent(totalLength); + char[] combinedArray = this.arrayPool.Rent(totalLength); // Copy chars from charsNotProcessedFromPreviousChunk to the combined array charsNotProcessedFromPreviousChunk.CopyTo(combinedArray); @@ -466,7 +509,7 @@ private async ValueTask WriteCharsInChunksWithoutEscapingAsync(ReadOnlyMemory.Shared.Return(combinedArray, true); + this.arrayPool.Return(combinedArray, true); } else { @@ -488,7 +531,7 @@ private void WriteChunk(ReadOnlySpan chunk, int chunkLength, int firstInde { if (this.buffer == null) { - this.buffer = ArrayPool.Shared.Rent(ODataUtf8JsonWriter.chunkSize); + this.buffer = this.arrayPool.Rent(ODataUtf8JsonWriter.chunkSize); } // Update the buffer with unprocessed bytes from the current chunk. @@ -522,7 +565,7 @@ private void WriteChunkWithoutEscaping(ReadOnlySpan chunk, int chunkLength { if (this.buffer == null) { - this.buffer = ArrayPool.Shared.Rent(ODataUtf8JsonWriter.chunkSize); + this.buffer = this.arrayPool.Rent(ODataUtf8JsonWriter.chunkSize); } // Update the buffer with unprocessed bytes from the current chunk. diff --git a/test/UnitTests/Microsoft.OData.Core.Tests/Json/MockJsonWriter.cs b/test/UnitTests/Microsoft.OData.Core.Tests/Json/MockJsonWriter.cs index 685b20265a..47c537c008 100644 --- a/test/UnitTests/Microsoft.OData.Core.Tests/Json/MockJsonWriter.cs +++ b/test/UnitTests/Microsoft.OData.Core.Tests/Json/MockJsonWriter.cs @@ -23,6 +23,10 @@ internal class MockJsonWriter : IJsonWriter { public Action WriteNameVerifier; public Action WriteValueVerifier; + public Func StartStreamValueScopeFunc; + public Action EndStreamValueScopeAction; + public Func> StartStreamValueScopeAsyncFunc; + public Func EndStreamValueScopeAsyncFunc; public void StartPaddingFunctionScope() => throw new NotImplementedException(); @@ -168,19 +172,19 @@ public Task WriteRawValueAsync(string rawValue) public Task WriteValueAsync(JsonElement value) => throw new NotImplementedException(); - public Stream StartStreamValueScope() => throw new NotImplementedException(); + public Stream StartStreamValueScope() => this.StartStreamValueScopeFunc?.Invoke() ?? throw new NotImplementedException(); public TextWriter StartTextWriterValueScope(string contentType) => throw new NotImplementedException(); - public void EndStreamValueScope() => throw new NotImplementedException(); + public void EndStreamValueScope() => (this.EndStreamValueScopeAction ?? throw new NotImplementedException()).Invoke(); public void EndTextWriterValueScope() => throw new NotImplementedException(); - public Task StartStreamValueScopeAsync() => throw new NotImplementedException(); + public Task StartStreamValueScopeAsync() => this.StartStreamValueScopeAsyncFunc?.Invoke() ?? throw new NotImplementedException(); public Task StartTextWriterValueScopeAsync(string contentType) => throw new NotImplementedException(); - public Task EndStreamValueScopeAsync() => throw new NotImplementedException(); + public Task EndStreamValueScopeAsync() => this.EndStreamValueScopeAsyncFunc?.Invoke() ?? throw new NotImplementedException(); public Task EndTextWriterValueScopeAsync() => throw new NotImplementedException(); } diff --git a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataJsonValueSerializerAsyncTests.cs b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataJsonValueSerializerAsyncTests.cs index e866963d94..7ecb92e10b 100644 --- a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataJsonValueSerializerAsyncTests.cs +++ b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataJsonValueSerializerAsyncTests.cs @@ -10,6 +10,8 @@ using System.Text; using System.Threading.Tasks; using Microsoft.OData.Core; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.OData.Core.Tests.DependencyInjection; using Microsoft.OData.Edm; using Microsoft.OData.Json; using Xunit; @@ -505,6 +507,24 @@ public async Task WriteStreamValueAsync_WritesStreamValue() Assert.Equal("\"CjEyMzQ1Njc4OTA=\"", result); } + [Fact] + public async Task WriteStreamValueAsync_DisposesWriterStreamOnce() + { + var writerStream = new DisposeTrackingMemoryStream(); + var jsonWriter = new MockJsonWriter + { + StartStreamValueScopeAsyncFunc = () => Task.FromResult(writerStream), + EndStreamValueScopeAsyncFunc = async () => await writerStream.DisposeAsync() + }; + var container = ServiceProviderHelper.BuildServiceProvider( + services => services.AddSingleton(new MockJsonWriterFactory(jsonWriter))); + var serializer = this.CreateODataJsonValueSerializer(true, container, true); + + await serializer.WriteStreamValueAsync(new ODataBinaryStreamValue(new MemoryStream(new byte[] { 1 }))); + + Assert.Equal(1, writerStream.DisposeCount); + } + private ODataJsonValueSerializer CreateODataJsonValueSerializer(bool writingResponse, IServiceProvider serviceProvider = null, bool isAsync = false) { var messageInfo = new ODataMessageInfo @@ -537,5 +557,20 @@ private async Task SetupJsonValueSerializerAndRunTestAsync(Func writerStream, + EndStreamValueScopeAction = () => writerStream.Dispose() + }; + var container = ServiceProviderHelper.BuildServiceProvider( + services => services.AddSingleton(new MockJsonWriterFactory(jsonWriter))); + var serializer = this.CreateODataJsonValueSerializer(true, container); + + serializer.WriteStreamValue(new ODataBinaryStreamValue(new MemoryStream(new byte[] { 1 }))); + + Assert.Equal(1, writerStream.DisposeCount); + } + private ODataJsonValueSerializer CreateODataJsonValueSerializer(bool writingResponse, IServiceProvider serviceProvider = null) { var messageInfo = new ODataMessageInfo @@ -398,11 +416,14 @@ private class DisposeTrackingMemoryStream : MemoryStream { public bool Disposed { get; private set; } = false; + public int DisposeCount { get; private set; } + protected override void Dispose(bool disposing) { if (disposing) { this.Disposed = true; + this.DisposeCount++; } base.Dispose(disposing); diff --git a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonTextWriterTests.cs b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonTextWriterTests.cs index 1ac35674c0..fa4b0a31e7 100644 --- a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonTextWriterTests.cs +++ b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonTextWriterTests.cs @@ -5,6 +5,10 @@ //--------------------------------------------------------------------- using System; +using System.IO; +using System.Text; +using System.Threading.Tasks; +using Microsoft.OData.Json; using static Microsoft.OData.Json.ODataUtf8JsonWriter; using Xunit; @@ -18,5 +22,21 @@ public void Encoding_ThrowsNotImplementedException() var stream = new ODataUtf8JsonTextWriter(null); Assert.Throws(() => stream.Encoding); } + + [Fact] + public async Task MixedDispose_ReturnsEachRentedBufferOnce() + { + var output = new MemoryStream(); + var jsonWriter = new ODataUtf8JsonWriter(output, false, Encoding.UTF8, leaveStreamOpen: true); + var arrayPool = new TrackingArrayPool(); + var textWriter = new ODataUtf8JsonTextWriter(jsonWriter, arrayPool); + + await textWriter.WriteAsync('a'); + await textWriter.DisposeAsync(); + textWriter.Dispose(); + + Assert.Equal(1, arrayPool.RentCount); + Assert.Equal(arrayPool.RentCount, arrayPool.ReturnCount); + } } } diff --git a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonWriterStreamTests.cs b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonWriterStreamTests.cs index d3206c9063..79d57dfdf3 100644 --- a/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonWriterStreamTests.cs +++ b/test/UnitTests/Microsoft.OData.Core.Tests/Json/ODataUtf8JsonWriterStreamTests.cs @@ -5,6 +5,12 @@ //--------------------------------------------------------------------- using System; +using System.Buffers; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using System.Threading.Tasks; +using Microsoft.OData.Json; using static Microsoft.OData.Json.ODataUtf8JsonWriter; using Xunit; using System.IO; @@ -76,5 +82,70 @@ public void SetLength_ThrowsNotSupportedException() var stream = new ODataUtf8JsonWriteStream(null); Assert.Throws(() => stream.SetLength(10)); } + + [Theory] + [InlineData(1, 1)] + [InlineData(2, 1)] + [InlineData(3, 0)] + [InlineData(2048, 1)] + [InlineData(2049, 2)] + public void DisposeTwice_ReturnsEachRentedBufferOnce(int payloadLength, int expectedRentCount) + { + var output = new MemoryStream(); + var jsonWriter = new ODataUtf8JsonWriter(output, false, Encoding.UTF8, leaveStreamOpen: true); + var arrayPool = new TrackingArrayPool(); + var stream = new ODataUtf8JsonWriteStream(jsonWriter, arrayPool); + byte[] payload = Enumerable.Range(0, payloadLength).Select(i => (byte)i).ToArray(); + + stream.Write(payload, 0, payload.Length); + stream.Dispose(); + stream.Dispose(); + jsonWriter.Flush(); + + Assert.Equal(Convert.ToBase64String(payload), Encoding.UTF8.GetString(output.ToArray())); + Assert.Equal(expectedRentCount, arrayPool.RentCount); + Assert.Equal(arrayPool.RentCount, arrayPool.ReturnCount); + } + + [Fact] + public async Task MixedDispose_ReturnsEachRentedBufferOnce() + { + var output = new MemoryStream(); + var jsonWriter = new ODataUtf8JsonWriter(output, false, Encoding.UTF8, leaveStreamOpen: true); + var arrayPool = new TrackingArrayPool(); + var stream = new ODataUtf8JsonWriteStream(jsonWriter, arrayPool); + + await stream.WriteAsync(new byte[2049], 0, 2049); + stream.Write(new byte[3], 0, 3); + await stream.DisposeAsync(); + stream.Dispose(); + await jsonWriter.FlushAsync(); + + Assert.Equal(2, arrayPool.RentCount); + Assert.Equal(arrayPool.RentCount, arrayPool.ReturnCount); + } + } + + internal sealed class TrackingArrayPool : ArrayPool + { + private readonly HashSet rentedArrays = new HashSet(); + private readonly HashSet returnedArrays = new HashSet(); + + public int RentCount => this.rentedArrays.Count; + + public int ReturnCount => this.returnedArrays.Count; + + public override T[] Rent(int minimumLength) + { + var array = new T[minimumLength]; + this.rentedArrays.Add(array); + return array; + } + + public override void Return(T[] array, bool clearArray = false) + { + Assert.Contains(array, this.rentedArrays); + Assert.True(this.returnedArrays.Add(array), "The same pooled array was returned more than once."); + } } }