System.ClientModel: Provide APIs for message types to transfer dispose ownership of response

sdk/core/Azure.Core/api/Azure.Core.net461.cs

@@ -496,6 +496,7 @@ public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public void SetProperty(System.Type type, object value) { }
+        public Azure.Response TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(string name, out object? value) { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/Azure.Core/api/Azure.Core.net472.cs

@@ -496,6 +496,7 @@ public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public void SetProperty(System.Type type, object value) { }
+        public Azure.Response TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(string name, out object? value) { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/Azure.Core/api/Azure.Core.net6.0.cs

@@ -496,6 +496,7 @@ public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public void SetProperty(System.Type type, object value) { }
+        public Azure.Response TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(string name, out object? value) { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs

@@ -496,6 +496,7 @@ public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
         public void SetProperty(System.Type type, object value) { }
+        public Azure.Response TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(string name, out object? value) { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

sdk/core/Azure.Core/src/HttpMessage.cs

@@ -16,6 +16,7 @@ public sealed class HttpMessage : IDisposable
     {
         private ArrayBackedPropertyBag<ulong, object> _propertyBag;
         private Response? _response;
+        private bool _ownsResponse;
 
         /// <summary>
         /// Creates a new instance of <see cref="HttpMessage"/>.
@@ -29,6 +30,7 @@ public HttpMessage(Request request, ResponseClassifier responseClassifier)
             ResponseClassifier = responseClassifier;
             BufferResponse = true;
             _propertyBag = new ArrayBackedPropertyBag<ulong, object>();
+            _ownsResponse = true;
         }
 
         /// <summary>
@@ -177,7 +179,9 @@ public void SetProperty(Type type, object value) =>
             _propertyBag.Set((ulong)type.TypeHandle.Value, value);
 
         /// <summary>
-        /// Returns the response content stream and releases it ownership to the caller. After calling this methods using <see cref="Azure.Response.ContentStream"/> or <see cref="Azure.Response.Content"/> would result in exception.
+        /// Returns the response content stream and releases it ownership to the caller.
+        /// After calling this methods using <see cref="Azure.Response.ContentStream"/>
+        /// or <see cref="Azure.Response.Content"/> would result in exception.
         /// </summary>
         /// <returns>The content stream or null if response didn't have any.</returns>
         public Stream? ExtractResponseContent()
@@ -194,6 +198,20 @@ public void SetProperty(Type type, object value) =>
             }
         }
 
+        /// <summary>
+        /// Transfers dispose ownership of the message's Response property to the
+        /// holder of the Response.  After calling this method, the <see cref="Response"/>
+        /// property will still hold the response object, but when this object is
+        /// disposed, it will no longer dispose the response as well.
+        /// </summary>
+        public Response TransferResponseDisposeOwnership()
+        {
+            Response response = Response;
+            _response = null;
+            _ownsResponse = false;
+            return response;
+        }
+
         /// <summary>
         /// Disposes the request and response.
         /// </summary>
@@ -202,11 +220,14 @@ public void Dispose()
             Request.Dispose();
             _propertyBag.Dispose();
 
-            var response = _response;
-            if (response != null)
+            if (_ownsResponse)
             {
-                _response = null;
-                response.Dispose();
+                var response = _response;
+                if (response != null)
+                {
+                    _response = null;
+                    response.Dispose();
+                }
             }
         }
 

sdk/core/Azure.Core/tests/HttpMessageTests.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Threading.Tasks;
+using System.Threading;
+using System;
 using Azure.Core.Pipeline;
 using Azure.Core.TestFramework;
 using NUnit.Framework;
@@ -380,7 +383,47 @@ public void AppliesHandlerWithLastSetWinsSemantics()
             Assert.IsTrue(message.ResponseClassifier.IsErrorResponse(message));
         }
 
+        [Test]
+        public async Task UnbufferedStreamAccessibleAfterMessageDisposed()
+        {
+            byte[] serverBytes = new byte[1000];
+            new Random().NextBytes(serverBytes);
+
+            HttpPipeline pipeline = HttpPipelineBuilder.Build(new MockClientOptions { Retry = { NetworkTimeout = Timeout.InfiniteTimeSpan } });
+
+            using TestServer testServer = new(async context =>
+            {
+                await context.Response.Body.WriteAsync(serverBytes, 0, serverBytes.Length).ConfigureAwait(false);
+            });
+
+            Response response;
+            using (HttpMessage message = pipeline.CreateMessage())
+            {
+                message.Request.Uri.Reset(testServer.Address);
+                message.BufferResponse = false;
+
+                await pipeline.SendAsync(message, default).ConfigureAwait(false);
+
+                response = message.TransferResponseDisposeOwnership();
+            }
+
+            Assert.NotNull(response.ContentStream);
+            byte[] clientBytes = new byte[serverBytes.Length];
+            int readLength = 0;
+            while (readLength < serverBytes.Length)
+            {
+                readLength += await response.ContentStream.ReadAsync(clientBytes, 0, serverBytes.Length);
+            }
+
+            Assert.AreEqual(serverBytes.Length, readLength);
+            CollectionAssert.AreEqual(serverBytes, clientBytes);
+        }
+
         #region Helpers
+        internal class MockClientOptions : ClientOptions
+        {
+        }
+
         private class StatusCodeHandler : ResponseClassificationHandler
         {
             private readonly int _statusCode;

sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs

@@ -150,6 +150,7 @@ public void Apply(System.ClientModel.Primitives.RequestOptions options) { }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
         public void SetProperty(System.Type type, object value) { }
+        public System.ClientModel.Primitives.PipelineResponse TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }
     public abstract partial class PipelineMessageClassifier

sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs

@@ -149,6 +149,7 @@ public void Apply(System.ClientModel.Primitives.RequestOptions options) { }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
         public void SetProperty(System.Type type, object value) { }
+        public System.ClientModel.Primitives.PipelineResponse TransferResponseDisposeOwnership() { throw null; }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }
     public abstract partial class PipelineMessageClassifier

sdk/core/System.ClientModel/src/Message/PipelineMessage.cs

@@ -10,6 +10,7 @@ public class PipelineMessage : IDisposable
 {
     private ArrayBackedPropertyBag<ulong, object> _propertyBag;
     private bool _disposed;
+    private bool _ownsResponse;
 
     protected internal PipelineMessage(PipelineRequest request)
     {
@@ -20,6 +21,8 @@ protected internal PipelineMessage(PipelineRequest request)
 
         BufferResponse = true;
         ResponseClassifier = PipelineMessageClassifier.Default;
+
+        _ownsResponse = true;
     }
 
     public PipelineRequest Request { get; }
@@ -125,33 +128,46 @@ internal void AssertResponse()
         }
     }
 
+    public PipelineResponse TransferResponseDisposeOwnership()
+    {
+        AssertResponse();
+
+        PipelineResponse response = Response!;
+        Response = null;
+        _ownsResponse = false;
+        return response;
+    }
+
     #region IDisposable
 
+    public void Dispose()
+    {
+        Dispose(true);
+        GC.SuppressFinalize(this);
+    }
+
     protected virtual void Dispose(bool disposing)
     {
         if (disposing && !_disposed)
         {
-            var request = Request;
+            PipelineRequest request = Request;
             request?.Dispose();
 
             _propertyBag.Dispose();
 
-            var response = Response;
-            if (response != null)
+            if (_ownsResponse)
             {
-                response.Dispose();
-                Response = null;
+                PipelineResponse? response = Response;
+                if (response != null)
+                {
+                    response.Dispose();
+                    Response = null;
+                }
             }
 
             _disposed = true;
         }
     }
 
-    public void Dispose()
-    {
-        Dispose(true);
-        GC.SuppressFinalize(this);
-    }
-
     #endregion
 }

sdk/core/System.ClientModel/tests/Message/PipelineMessageTests.cs

@@ -1,14 +1,22 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using NUnit.Framework;
 using System.ClientModel.Primitives;
 using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core.TestFramework;
+using ClientModel.Tests.Mocks;
+using NUnit.Framework;
+using SyncAsyncTestBase = ClientModel.Tests.SyncAsyncTestBase;
 
 namespace System.ClientModel.Tests.Message;
 
-public class PipelineMessageTests
+public class PipelineMessageTests : SyncAsyncTestBase
 {
+    public PipelineMessageTests(bool isAsync) : base(isAsync)
+    {
+    }
+
     [Test]
     public void ApplyAddsRequestHeaders()
     {
@@ -111,6 +119,43 @@ public void TryGetTypeKeyedPropertyReturnsCorrectValues()
         Assert.AreEqual(4444, ((T4)value!).Value);
     }
 
+    [Test]
+    public async Task UnbufferedStreamAccessibleAfterMessageDisposed()
+    {
+        byte[] serverBytes = new byte[1000];
+        new Random().NextBytes(serverBytes);
+
+        ClientPipelineOptions options = new() { NetworkTimeout = Timeout.InfiniteTimeSpan };
+        ClientPipeline pipeline = ClientPipeline.Create(options);
+
+        using TestServer testServer = new(async context =>
+        {
+            await context.Response.Body.WriteAsync(serverBytes, 0, serverBytes.Length).ConfigureAwait(false);
+        });
+
+        PipelineResponse response;
+        using (PipelineMessage message = pipeline.CreateMessage())
+        {
+            message.Request.Uri = testServer.Address;
+            message.BufferResponse = false;
+
+            await pipeline.SendSyncOrAsync(message, IsAsync);
+
+            response = message.TransferResponseDisposeOwnership();
+        }
+
+        Assert.NotNull(response.ContentStream);
+        byte[] clientBytes = new byte[serverBytes.Length];
+        int readLength = 0;
+        while (readLength < serverBytes.Length)
+        {
+            readLength += await response.ContentStream!.ReadAsync(clientBytes, 0, serverBytes.Length);
+        }
+
+        Assert.AreEqual(serverBytes.Length, readLength);
+        CollectionAssert.AreEqual(serverBytes, clientBytes);
+    }
+
     #region Helpers
     private struct T1
     {