Fix connection leak in rest proxy when return type is void or mono<vo…

…id> (Azure#30072) * Fix connection leak in rest proxy when return type is void or mono<void> * changelog.
azure-sdk · Jul 21, 2022 · 92b7424 · 92b7424
1 parent 3f0138c
commit 92b7424
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 20 deletions.
diff --git a/sdk/core/azure-core-perf/src/main/java/com/azure/core/perf/BinaryDataSendTest.java b/sdk/core/azure-core-perf/src/main/java/com/azure/core/perf/BinaryDataSendTest.java
@@ -28,7 +28,7 @@ public void run() {
 
     @Override
     public Mono<Void> runAsync() {
-        return service.setBinaryData(endpoint, id, binaryDataSupplier.get(), length)
-            .then();
+        return Mono.fromSupplier(binaryDataSupplier)
+            .flatMap(data -> service.setBinaryData(endpoint, id, data, length));
     }
 }
diff --git a/sdk/core/azure-core-perf/src/main/java/com/azure/core/perf/PipelineSendTest.java b/sdk/core/azure-core-perf/src/main/java/com/azure/core/perf/PipelineSendTest.java
@@ -48,20 +48,23 @@ public void run() {
 
     @Override
     public Mono<Void> runAsync() {
-        HttpHeaders headers = new HttpHeaders();
-        headers.set("Content-Length", contentLengthHeaderValue);
-        HttpRequest httpRequest = new HttpRequest(
-            HttpMethod.PUT, targetURL, headers, binaryDataSupplier.get().toFluxByteBuffer());
-        // Context with azure-eagerly-read-response=true makes sure
-        // that response is disposed to prevent connection leak.
-        // There's no response body in this scenario anyway.
-        return httpPipeline.send(httpRequest, context)
-            .map(httpResponse -> {
-                if (httpResponse.getStatusCode() / 100 != 2) {
-                    throw new IllegalStateException("Endpoint didn't return 2xx http status code.");
-                }
-                return httpResponse;
-            })
-            .then();
+        return Mono.fromSupplier(binaryDataSupplier)
+            .flatMap(data -> {
+                HttpHeaders headers = new HttpHeaders();
+                headers.set("Content-Length", contentLengthHeaderValue);
+                HttpRequest httpRequest = new HttpRequest(
+                    HttpMethod.PUT, targetURL, headers, data);
+                // Context with azure-eagerly-read-response=true makes sure
+                // that response is disposed to prevent connection leak.
+                // There's no response body in this scenario anyway.
+                return httpPipeline.send(httpRequest, context)
+                    .map(httpResponse -> {
+                        if (httpResponse.getStatusCode() / 100 != 2) {
+                            throw new IllegalStateException("Endpoint didn't return 2xx http status code.");
+                        }
+                        return httpResponse;
+                    })
+                    .then();
+            });
     }
 }
diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed bug where `RestProxy` could leak connection if service method returned `Mono<Void>` or `void`.
+
 ### Other Changes
 
 ## 1.30.0 (2022-06-30)

diff --git a/...core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java b/...core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java
@@ -206,7 +206,7 @@ private Object handleRestReturnType(final Mono<HttpResponseDecoder.HttpDecodedRe
             final Type monoTypeParam = TypeUtil.getTypeArgument(returnType);
             if (TypeUtil.isTypeOrSubTypeOf(monoTypeParam, Void.class)) {
                 // ProxyMethod ReturnType: Mono<Void>
-                result = asyncExpectedResponse.then();
+                result = asyncExpectedResponse.doOnNext(HttpResponseDecoder.HttpDecodedResponse::close).then();
             } else {
                 // ProxyMethod ReturnType: Mono<? extends RestResponseBase<?, ?>>
                 result = asyncExpectedResponse.flatMap(response ->
@@ -218,7 +218,7 @@ private Object handleRestReturnType(final Mono<HttpResponseDecoder.HttpDecodedRe
         } else if (TypeUtil.isTypeOrSubTypeOf(returnType, void.class) || TypeUtil.isTypeOrSubTypeOf(returnType,
             Void.class)) {
             // ProxyMethod ReturnType: Void
-            asyncExpectedResponse.block();
+            asyncExpectedResponse.doOnNext(HttpResponseDecoder.HttpDecodedResponse::close).block();
             result = null;
         } else {
             // ProxyMethod ReturnType: T where T != async (Mono, Flux) or sync Void

diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java
@@ -210,7 +210,8 @@ private Object handleRestReturnType(final HttpResponseDecoder.HttpDecodedRespons
         if (TypeUtil.isTypeOrSubTypeOf(returnType, void.class) || TypeUtil.isTypeOrSubTypeOf(returnType,
             Void.class)) {
             // ProxyMethod ReturnType: Void
-            result = expectedResponse;
+            expectedResponse.close();
+            result = null;
         } else {
             // ProxyMethod ReturnType: T where T != async (Mono, Flux) or sync Void
             // Block the deserialization until a value T is received

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
@@ -70,6 +70,14 @@ Mono<Response<Void>> testMethod(
             @HeaderParam("Content-Length") Long contentLength
         );
 
+        @Get("my/url/path")
+        @ExpectedResponses({200})
+        Mono<Void> testMethodReturnsMonoVoid();
+
+        @Get("my/url/path")
+        @ExpectedResponses({200})
+        void testVoidMethod();
+
         @Get("my/url/path")
         @ExpectedResponses({200})
         StreamResponse testDownload();
@@ -179,6 +187,35 @@ public void doesNotChangeBinaryDataContentType(BinaryData data, long contentLeng
         assertEquals(expectedContentClazz, actualContentClazz);
     }
 
+    @Test
+    public void monoVoidReturningApiClosesResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+        StepVerifier.create(
+                testInterface.testMethodReturnsMonoVoid())
+            .verifyComplete();
+
+        Mockito.verify(client.lastResponseSpy).close();
+    }
+
+    @Test
+    public void voidReturningApiClosesResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+
+        testInterface.testVoidMethod();
+
+        Mockito.verify(client.lastResponseSpy).close();
+    }
+
     private static Stream<Arguments> doesNotChangeBinaryDataContentTypeDataProvider() throws Exception {
         String string = "hello";
         byte[] bytes = string.getBytes();

diff --git a/.../azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java b/.../azure-core/src/test/java/com/azure/core/implementation/http/rest/SyncRestProxyTests.java
@@ -62,6 +62,25 @@ Response<Void> testMethod(
         @Get("my/url/path")
         @ExpectedResponses({200})
         StreamResponse testDownload(Context context);
+
+        @Get("my/url/path")
+        @ExpectedResponses({200})
+        void testVoidMethod(Context context);
+    }
+
+    @Test
+    public void voidReturningApiClosesResponse() {
+        LocalHttpClient client = new LocalHttpClient();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(client)
+            .build();
+
+        TestInterface testInterface = RestProxy.create(TestInterface.class, pipeline);
+
+        Context context =  new Context(HTTP_REST_PROXY_SYNC_PROXY_ENABLE, true);
+        testInterface.testVoidMethod(context);
+
+        Mockito.verify(client.lastResponseSpy).close();
     }
 
     @Test