Code changes to use cloneable stream in buffered stream.

Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs

@@ -108,11 +108,9 @@ public override async Task<ResponseMessage> SendAsync(
             if (ConfigurationManager.IsBinaryEncodingEnabled()
                 && RequestInvokerHandler.IsPointOperationSupportedForBinaryEncoding(request)
                 && response.Content != null
-                && (response.Content is not MemoryStream 
-                || response.Content is not CloneableStream))
+                && response.Content is not CloneableStream)
             {
-                CloneableStream clonableStream = await StreamExtension.AsClonableStreamAsync(response.Content, default);
-                response.Content = clonableStream;
+                response.Content = await StreamExtension.AsClonableStreamAsync(response.Content, default);
             }
 
             return response;

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs

@@ -907,7 +907,7 @@ private async Task<ResponseMessage> ProcessItemStreamAsync(
             OperationType operationType,
             ItemRequestOptions requestOptions,
             ITrace trace,
-            JsonSerializationFormat? targetResponseSerializationFormat,
+            JsonSerializationFormat targetResponseSerializationFormat,
             CancellationToken cancellationToken)
         {
             if (trace == null)
@@ -924,14 +924,9 @@ private async Task<ResponseMessage> ProcessItemStreamAsync(
             string resourceUri = this.GetResourceUri(requestOptions, operationType, itemId);
 
             // Convert Text to Binary Stream.
-            Stream convertedStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(
+            streamPayload = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(
                 targetSerializationFormat: ContainerCore.GetTargetRequestSerializationFormat(),
-                inputStream: streamPayload);
-
-            if (convertedStream != null)
-            {
-                streamPayload = convertedStream;
-            }
+                inputStream: streamPayload == null ? null : await StreamExtension.AsClonableStreamAsync(streamPayload));
 
             ResponseMessage responseMessage = await this.ClientContext.ProcessResourceOperationStreamAsync(
                 resourceUri: resourceUri,
@@ -947,16 +942,12 @@ private async Task<ResponseMessage> ProcessItemStreamAsync(
                 cancellationToken: cancellationToken);
 
             // Convert Binary Stream to Text.
-            if (requestOptions == null || !requestOptions.EnableBinaryResponseOnPointOperations)
+            if ((requestOptions == null || !requestOptions.EnableBinaryResponseOnPointOperations)
+                && responseMessage?.Content is CloneableStream outputCloneableStream)
             {
-                Stream outputResponseStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(
+                responseMessage.Content = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(
                     targetSerializationFormat: targetResponseSerializationFormat,
-                    inputStream: responseMessage?.Content);
-
-                if (outputResponseStream != null)
-                {
-                    responseMessage.Content = outputResponseStream;
-                }
+                    inputStream: outputCloneableStream);
             }
 
             return responseMessage;

Microsoft.Azure.Cosmos/src/Serializer/CosmosBufferedStreamWrapper.cs

@@ -5,7 +5,6 @@
 namespace Microsoft.Azure.Cosmos.Serializer
 {
     using System;
-    using System.Diagnostics;
     using System.IO;
     using System.Threading;
     using System.Threading.Tasks;
@@ -20,7 +19,7 @@ internal class CosmosBufferedStreamWrapper : Stream
         /// <summary>
         /// The inner stream being wrapped.
         /// </summary>
-        private readonly Stream innerStream;
+        private readonly CloneableStream innerStream;
 
         /// <summary>
         /// Indicates whether the inner stream should be disposed.
@@ -43,13 +42,9 @@ internal class CosmosBufferedStreamWrapper : Stream
         /// <param name="inputStream">The input stream to wrap.</param>
         /// <param name="shouldDisposeInnerStream">Indicates whether the inner stream should be disposed.</param>
         public CosmosBufferedStreamWrapper(
-            Stream inputStream,
+            CloneableStream inputStream,
             bool shouldDisposeInnerStream)
         {
-            Debug.Assert(
-                inputStream is CloneableStream || inputStream is MemoryStream,
-                "The inner stream is neither a memory stream nor a cloneable stream.");
-
             this.innerStream = inputStream ?? throw new ArgumentNullException(nameof(inputStream));
             this.shouldDisposeInnerStream = shouldDisposeInnerStream;
         }
@@ -99,33 +94,7 @@ public override int Read(byte[] buffer, int offset, int count)
                 throw new ArgumentNullException(nameof(buffer));
             }
 
-            if (offset < 0
-                || count < 0
-                || (buffer.Length - offset) < count
-                || this.innerStream.Position == this.innerStream.Length)
-            {
-                return 0;
-            }
-
-            int bytesRead = 0;
-            if (this.hasReadFirstByte
-                && this.innerStream.Position == 1
-                && offset == 0
-                && count > 0)
-            {
-                buffer[0] = this.firstByteBuffer[0];
-                bytesRead = 1;
-                offset++;
-                count--;
-            }
-
-            if (count > 0)
-            {
-                int innerBytesRead = this.innerStream.Read(buffer, offset, count);
-                bytesRead += innerBytesRead;
-            }
-
-            return bytesRead;
+            return this.innerStream.Read(buffer, offset, count);
         }
 
         /// <inheritdoc />
@@ -146,10 +115,7 @@ protected override void Dispose(bool disposing)
                 }
                 else
                 {
-                    if (this.innerStream.CanSeek)
-                    {
-                        this.innerStream.Position = 0;
-                    }
+                    this.ResetStreamPosition();
                 }
             }
 
@@ -173,12 +139,6 @@ public byte[] ReadAll()
                 totalBytes += count;
             }
 
-            if (this.hasReadFirstByte)
-            {
-                bytes[0] = this.firstByteBuffer[0];
-                totalBytes += 1;
-            }
-
             return totalBytes > 0 ? bytes : default;
         }
 
@@ -199,12 +159,6 @@ public async Task<byte[]> ReadAllAsync(CancellationToken cancellationToken = def
                 totalBytes += count;
             }
 
-            if (this.hasReadFirstByte)
-            {
-                bytes[0] = this.firstByteBuffer[0];
-                totalBytes += 1;
-            }
-
             return totalBytes > 0 ? bytes : default;
         }
 
@@ -216,7 +170,7 @@ public async Task<byte[]> ReadAllAsync(CancellationToken cancellationToken = def
         /// </returns>
         public JsonSerializationFormat GetJsonSerializationFormat()
         {
-            this.ReadFirstByte();
+            this.ReadFirstByteAndResetStream();
             if (this.firstByteBuffer[0] == (byte)JsonSerializationFormat.Binary)
             {
                 return JsonSerializationFormat.Binary;
@@ -230,17 +184,29 @@ public JsonSerializationFormat GetJsonSerializationFormat()
         }
 
         /// <summary>
-        /// Reads the first byte from the inner stream and stores it in the buffer.
+        /// Reads the first byte from the inner stream and stores it in the buffer. It also resets the stream position to zero.
         /// </summary>
         /// <remarks>
         /// This method sets the <see cref="hasReadFirstByte"/> flag to true if the first byte is successfully read.
         /// </remarks>
-        private void ReadFirstByte()
+        private void ReadFirstByteAndResetStream()
         {
             if (!this.hasReadFirstByte
                 && this.innerStream.Read(this.firstByteBuffer, 0, 1) > 0)
             {
                 this.hasReadFirstByte = true;
+                this.ResetStreamPosition();
+            }
+        }
+
+        /// <summary>
+        /// Resets the inner stream position to zero.
+        /// </summary>
+        private void ResetStreamPosition()
+        {
+            if (this.innerStream.CanSeek)
+            {
+                this.innerStream.Position = 0;
             }
         }
     }

Microsoft.Azure.Cosmos/src/Serializer/CosmosJsonDotNetSerializer.cs

@@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos
     using System.IO;
     using System.Text;
     using Microsoft.Azure.Cosmos.Serializer;
+    using Microsoft.Azure.Documents;
     using Newtonsoft.Json;
     using Newtonsoft.Json.Serialization;
 
@@ -88,7 +89,7 @@ internal CosmosJsonDotNetSerializer(
         /// <returns>The object representing the deserialized stream</returns>
         public override T FromStream<T>(Stream stream)
         {
-            using (CosmosBufferedStreamWrapper bufferedStream = new (stream, shouldDisposeInnerStream: true))
+            using (stream)
             {
                 if (typeof(Stream).IsAssignableFrom(typeof(T)))
                 {
@@ -97,25 +98,29 @@ public override T FromStream<T>(Stream stream)
 
                 JsonSerializer jsonSerializer = this.GetSerializer();
 
-                if (bufferedStream.GetJsonSerializationFormat() == Json.JsonSerializationFormat.Binary)
+                if (stream is CloneableStream cloneableStream)
                 {
-                    byte[] content = bufferedStream.ReadAll();
+                    using (CosmosBufferedStreamWrapper bufferedStream = new (cloneableStream, shouldDisposeInnerStream: false))
+                    {
+                        if (bufferedStream.GetJsonSerializationFormat() == Json.JsonSerializationFormat.Binary)
+                        {
+                            byte[] content = bufferedStream.ReadAll();
 
-                    using Json.Interop.CosmosDBToNewtonsoftReader reader = new (
-                        jsonReader: Json.JsonReader.Create(
-                            jsonSerializationFormat: Json.JsonSerializationFormat.Binary,
-                            buffer: content));
+                            using Json.Interop.CosmosDBToNewtonsoftReader reader = new (
+                                jsonReader: Json.JsonReader.Create(
+                                    jsonSerializationFormat: Json.JsonSerializationFormat.Binary,
+                                    buffer: content));
 
-                    return jsonSerializer.Deserialize<T>(reader);
+                            return jsonSerializer.Deserialize<T>(reader);
+                        }
+                    }
                 }
-                else
+
+                using (StreamReader sr = new (stream))
                 {
-                    using (StreamReader sr = new (bufferedStream))
+                    using (JsonTextReader jsonTextReader = new (sr))
                     {
-                        using (JsonTextReader jsonTextReader = new (sr))
-                        {
-                            return jsonSerializer.Deserialize<T>(jsonTextReader);
-                        }
+                        return jsonSerializer.Deserialize<T>(jsonTextReader);
                     }
                 }
             }

Microsoft.Azure.Cosmos/src/Serializer/CosmosSerializationUtil.cs

@@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.Json;
     using Microsoft.Azure.Cosmos.Json.Interop;
     using Microsoft.Azure.Cosmos.Serializer;
+    using Microsoft.Azure.Documents;
     using Newtonsoft.Json.Serialization;
 
     internal static class CosmosSerializationUtil
@@ -45,10 +46,10 @@ internal static string GetStringWithPropertyNamingPolicy(CosmosPropertyNamingPol
         /// <param name="inputStream">The input stream containing the data to be serialized.</param>
         /// <returns>Returns true if the input stream is successfully serialized to the target format, otherwise false.</returns>
         internal static async Task<Stream> TrySerializeStreamToTargetFormatAsync(
-            JsonSerializationFormat? targetSerializationFormat,
-            Stream inputStream)
+            JsonSerializationFormat targetSerializationFormat,
+            CloneableStream inputStream)
         {
-            if (targetSerializationFormat == null || inputStream == null)
+            if (inputStream == null)
             {
                 return null;
             }
@@ -68,12 +69,12 @@ internal static async Task<Stream> TrySerializeStreamToTargetFormatAsync(
                     {
                         return CosmosSerializationUtil.ConvertToStreamUsingJsonSerializationFormat(
                             targetContent,
-                            targetSerializationFormat.Value);
+                            targetSerializationFormat);
                     }
                 }
             }
 
-            return null;
+            return inputStream;
         }
 
         /// <summary>

Microsoft.Azure.Cosmos/src/Serializer/CosmosSystemTextJsonSerializer.cs

@@ -50,24 +50,30 @@ public override T FromStream<T>(Stream stream)
                 return default;
             }
 
-            using (CosmosBufferedStreamWrapper bufferedStream = new (stream, shouldDisposeInnerStream: true))
+            using (stream)
             {
-                if (bufferedStream.GetJsonSerializationFormat() == Json.JsonSerializationFormat.Binary)
+                if (stream is Documents.CloneableStream cloneableStream)
                 {
-                    byte[] content = bufferedStream.ReadAll();
-
-                    if (CosmosObject.TryCreateFromBuffer(content, out CosmosObject cosmosObject))
-                    {
-                        return System.Text.Json.JsonSerializer.Deserialize<T>(cosmosObject.ToString(), this.jsonSerializerOptions);
-                    }
-                    else
+                    using (CosmosBufferedStreamWrapper bufferedStream = new (cloneableStream, shouldDisposeInnerStream: false))
                     {
-                        using Stream textStream = CosmosSerializationUtil.ConvertToStreamUsingJsonSerializationFormat(content, JsonSerializationFormat.Text);
-                        return this.DeserializeStream<T>(textStream);
+                        if (bufferedStream.GetJsonSerializationFormat() == JsonSerializationFormat.Binary)
+                        {
+                            byte[] content = bufferedStream.ReadAll();
+
+                            if (CosmosObject.TryCreateFromBuffer(content, out CosmosObject cosmosObject))
+                            {
+                                return System.Text.Json.JsonSerializer.Deserialize<T>(cosmosObject.ToString(), this.jsonSerializerOptions);
+                            }
+                            else
+                            {
+                                using Stream textStream = CosmosSerializationUtil.ConvertToStreamUsingJsonSerializationFormat(content, JsonSerializationFormat.Text);
+                                return this.DeserializeStream<T>(textStream);
+                            }
+                        }
                     }
                 }
 
-                return this.DeserializeStream<T>(bufferedStream);
+                return this.DeserializeStream<T>(stream);
             }
         }
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosSerializationUtilTests.cs

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Tests
     using System.IO;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Json;
+    using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
 
     /// <summary>
@@ -114,7 +115,8 @@ public async Task TrySerializeStreamToTargetFormat_Success(string expected, stri
                 : CosmosSerializerUtils.ConvertInputToBinaryStream(json, Newtonsoft.Json.JsonSerializer.Create());
 
             // Act
-            Stream outputStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(targetFormat, inputStream);
+            CloneableStream cloneableStream = await StreamExtension.AsClonableStreamAsync(inputStream);
+            Stream outputStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(targetFormat, cloneableStream);
 
             // Assert
             Assert.IsNotNull(outputStream);
@@ -132,7 +134,8 @@ public async Task TrySerializeStreamToTargetFormat_Failure()
             JsonSerializationFormat targetFormat = JsonSerializationFormat.Text;
 
             // Act
-            Stream outputStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(targetFormat, inputStream);
+            CloneableStream cloneableStream = await StreamExtension.AsClonableStreamAsync(inputStream);
+            Stream outputStream = await CosmosSerializationUtil.TrySerializeStreamToTargetFormatAsync(targetFormat, cloneableStream);
 
             // Assert
             Assert.IsNull(outputStream);