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

sdk/core/System.ClientModel/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added `ExtractResponse` method to `PipelineMessage` to enable returning an undisposed `PipelineResponse` from protocol methods.
+
 ### Breaking Changes
 
 - Change `HttpClientPipelineTransport.Shared` from a field to a property.

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

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

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

@@ -148,6 +148,7 @@ protected internal PipelineMessage(System.ClientModel.Primitives.PipelineRequest
         public void Apply(System.ClientModel.Primitives.RequestOptions options) { }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
+        public System.ClientModel.Primitives.PipelineResponse? ExtractResponse() { throw null; }
         public void SetProperty(System.Type type, object value) { }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }

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

@@ -32,6 +32,21 @@ public PipelineResponse? Response
         protected internal set;
     }
 
+    public PipelineResponse? ExtractResponse()
+    {
+        PipelineResponse? response = Response;
+        Response = null;
+        return response;
+    }
+
+    internal void AssertResponse()
+    {
+        if (Response is null)
+        {
+            throw new InvalidOperationException($"'{nameof(Response)}' property is not set on message.");
+        }
+    }
+
     #region Pipeline invocation options
 
     public CancellationToken CancellationToken
@@ -117,41 +132,30 @@ PerTryPolicies is not null ||
 
     #endregion
 
-    internal void AssertResponse()
+    #region IDisposable
+
+    public void Dispose()
     {
-        if (Response is null)
-        {
-            throw new InvalidOperationException("Response is not set on message.");
-        }
+        Dispose(true);
+        GC.SuppressFinalize(this);
     }
 
-    #region IDisposable
-
     protected virtual void Dispose(bool disposing)
     {
         if (disposing && !_disposed)
         {
-            var request = Request;
+            PipelineResponse? response = Response;
+            response?.Dispose();
+            Response = null;
+
+            PipelineRequest request = Request;
             request?.Dispose();
 
             _propertyBag.Dispose();
 
-            var 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,48 @@ public void TryGetTypeKeyedPropertyReturnsCorrectValues()
         Assert.AreEqual(4444, ((T4)value!).Value);
     }
 
+    [Test]
+    [TestCase(true)]
+    [TestCase(false)]
+    public async Task ResponseStreamAccessibleAfterMessageDisposed(bool buffer)
+    {
+        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 = buffer;
+
+            await pipeline.SendSyncOrAsync(message, IsAsync);
+
+            response = message.ExtractResponse();
+
+            Assert.IsNull(message.Response);
+        }
+
+        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
     {