Performance: Fix MemoryStream with un-accesible buffer (#1515)

* memorystream with buffer * test
Azure · May 14, 2020 · 01e39d4 · 01e39d4
1 parent 07ae4ce
commit 01e39d4
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 5 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs
@@ -148,11 +148,12 @@ private async Task<ResponseMessage> ReadNextInternalAsync(
             MemoryStream rewrittenMemoryStream;
             if (MemoryMarshal.TryGetArray(result, out ArraySegment<byte> rewrittenSegment))
             {
-                rewrittenMemoryStream = new MemoryStream(rewrittenSegment.Array, index: rewrittenSegment.Offset, count: rewrittenSegment.Count);
+                rewrittenMemoryStream = new MemoryStream(rewrittenSegment.Array, index: rewrittenSegment.Offset, count: rewrittenSegment.Count, writable: false, publiclyVisible: true);
             }
             else
             {
-                rewrittenMemoryStream = new MemoryStream(result.ToArray());
+                byte[] toArray = result.ToArray();
+                rewrittenMemoryStream = new MemoryStream(toArray, index: 0, count: toArray.Length, writable: false, publiclyVisible: true);
             }
 
             responseMessage.Content = rewrittenMemoryStream;

diff --git a/Microsoft.Azure.Cosmos/src/FeedRangeIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedRangeIteratorCore.cs
@@ -219,11 +219,12 @@ private async Task<ResponseMessage> ReadNextInternalAsync(
                 MemoryStream rewrittenMemoryStream;
                 if (MemoryMarshal.TryGetArray(result, out ArraySegment<byte> rewrittenSegment))
                 {
-                    rewrittenMemoryStream = new MemoryStream(rewrittenSegment.Array, index: rewrittenSegment.Offset, count: rewrittenSegment.Count);
+                    rewrittenMemoryStream = new MemoryStream(rewrittenSegment.Array, index: rewrittenSegment.Offset, count: rewrittenSegment.Count, writable: false, publiclyVisible: true);
                 }
                 else
                 {
-                    rewrittenMemoryStream = new MemoryStream(result.ToArray());
+                    byte[] toArray = result.ToArray();
+                    rewrittenMemoryStream = new MemoryStream(toArray, index: 0, count: toArray.Length, writable: false, publiclyVisible: true);
                 }
 
                 responseMessage.Content = rewrittenMemoryStream;

diff --git a/Microsoft.Azure.Cosmos/src/Serializer/CosmosElementSerializer.cs b/Microsoft.Azure.Cosmos/src/Serializer/CosmosElementSerializer.cs
@@ -242,7 +242,7 @@ internal static MemoryStream ToStream(
                 resultAsArray = new ArraySegment<byte>(result.ToArray());
             }
 
-            return new MemoryStream(resultAsArray.Array, resultAsArray.Offset, resultAsArray.Count);
+            return new MemoryStream(resultAsArray.Array, resultAsArray.Offset, resultAsArray.Count, writable: false, publiclyVisible: true);
         }
 
         internal static IEnumerable<T> GetResources<T>(

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs
@@ -94,6 +94,24 @@ public async Task MemoryStreamBufferIsAccessibleOnResponse()
             ResponseMessage response = await this.Container.CreateItemStreamAsync(streamPayload: TestCommon.SerializerCore.ToStream(testItem), partitionKey: new Cosmos.PartitionKey(testItem.status));
             Assert.IsNotNull(response);
             Assert.IsTrue((response.Content as MemoryStream).TryGetBuffer(out _));
+            FeedIterator feedIteratorQuery = this.Container.GetItemQueryStreamIterator(queryText: "SELECT * FROM c");
+
+            while (feedIteratorQuery.HasMoreResults)
+            {
+                ResponseMessage feedResponseQuery = await feedIteratorQuery.ReadNextAsync();
+                Assert.IsTrue((feedResponseQuery.Content as MemoryStream).TryGetBuffer(out _));
+            }
+
+            FeedIterator feedIterator = this.Container.GetItemQueryStreamIterator(requestOptions: new QueryRequestOptions()
+            {
+                PartitionKey = new Cosmos.PartitionKey(testItem.status)
+            });
+
+            while (feedIterator.HasMoreResults)
+            {
+                ResponseMessage feedResponse = await feedIterator.ReadNextAsync();
+                Assert.IsTrue((feedResponse.Content as MemoryStream).TryGetBuffer(out _));
+            }
         }
 
         [TestMethod]