Eagerly read response on Mono<Void> (Azure#30086)

* Fix connection leak in rest proxy when return type is void or mono<void> * changelog. * fix buffering for Mono<Void> * PR feedback. * tests for streamresponse.
azure-sdk · Jul 21, 2022 · 3cea40a · 3cea40a
1 parent 92b7424
commit 3cea40a
Show file tree

Hide file tree

Showing 2 changed files with 125 additions and 3 deletions.
diff --git a/...-core/src/main/java/com/azure/core/implementation/serializer/HttpResponseBodyDecoder.java b/...-core/src/main/java/com/azure/core/implementation/serializer/HttpResponseBodyDecoder.java
@@ -369,9 +369,14 @@ public static boolean shouldEagerlyReadResponse(Type returnType) {
             return false;
         }
 
-        return isReturnTypeDecodable(returnType)
-            || TypeUtil.isTypeOrSubTypeOf(returnType, Void.TYPE)
-            || TypeUtil.isTypeOrSubTypeOf(returnType, Void.class);
+
+        return isReturnTypeDecodable(returnType) || doesNotNeedResponseBody(returnType);
+    }
+
+    private static boolean doesNotNeedResponseBody(Type type) {
+        Type typeToInspect = unwrapReturnType(type);
+        return TypeUtil.isTypeOrSubTypeOf(typeToInspect, Void.TYPE)
+            || TypeUtil.isTypeOrSubTypeOf(typeToInspect, Void.class);
     }
 
     private static Type unwrapReturnType(Type returnType) {

diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/rest/RestProxyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/rest/RestProxyTests.java
@@ -43,6 +43,7 @@
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -78,6 +79,14 @@ Mono<Response<Void>> testMethod(
         @ExpectedResponses({200})
         void testVoidMethod();
 
+        @Get("my/url/path")
+        @ExpectedResponses({200})
+        Mono<Response<Void>> testMethodReturnsMonoResponseVoid();
+
+        @Get("my/url/path")
+        @ExpectedResponses({200})
+        Response<Void> testMethodReturnsResponseVoid();
+
         @Get("my/url/path")
         @ExpectedResponses({200})
         StreamResponse testDownload();
@@ -216,6 +225,103 @@ public void voidReturningApiClosesResponse() {
         Mockito.verify(client.lastResponseSpy).close();
     }
 
+    @Test
+    public void voidReturningApiEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+
+        testInterface.testVoidMethod();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertTrue((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
+    @Test
+    public void monoVoidReturningApiEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        StepVerifier.create(
+                testInterface.testMethodReturnsMonoVoid())
+            .verifyComplete();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertTrue((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
+    @Test
+    public void monoResponseVoidReturningApiEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        StepVerifier.create(
+                testInterface.testMethodReturnsMonoResponseVoid())
+            .expectNextCount(1)
+            .verifyComplete();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertTrue((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
+    @Test
+    public void responseVoidReturningApiEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        testInterface.testMethodReturnsResponseVoid();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertTrue((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
+    @Test
+    public void streamResponseDoesNotEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        testInterface.testDownload();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertFalse((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
+    @Test
+    public void monoWithStreamResponseDoesNotEagerlyReadsResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        StepVerifier.create(
+                testInterface.testDownloadAsync()
+                    .doOnNext(StreamResponse::close))
+            .expectNextCount(1)
+            .verifyComplete();
+
+        assertTrue(client.lastContext.getData("azure-eagerly-read-response").isPresent());
+        assertFalse((Boolean) client.lastContext.getData("azure-eagerly-read-response").get());
+    }
+
     private static Stream<Arguments> doesNotChangeBinaryDataContentTypeDataProvider() throws Exception {
         String string = "hello";
         byte[] bytes = string.getBytes();
@@ -243,10 +349,17 @@ private static final class LocalHttpClient implements HttpClient {
 
         private volatile HttpRequest lastHttpRequest;
         private volatile HttpResponse lastResponseSpy;
+        private volatile Context lastContext;
 
         @Override
         public Mono<HttpResponse> send(HttpRequest request) {
+            return send(request, Context.NONE);
+        }
+
+        @Override
+        public Mono<HttpResponse> send(HttpRequest request, Context context) {
             lastHttpRequest = request;
+            lastContext = context;
             boolean success = request.getUrl().getPath().equals("/my/url/path");
             if (request.getHttpMethod().equals(HttpMethod.POST)) {
                 success = success && request.getHeaders()
@@ -270,6 +383,10 @@ public HttpRequest getLastHttpRequest() {
         public HttpResponse getLastResponseSpy() {
             return lastResponseSpy;
         }
+
+        public Context getLastContext() {
+            return lastContext;
+        }
     }
 
     @ParameterizedTest