Fix blocking call in AsyncRestProxy (#36384)

Azure · Aug 18, 2023 · 317aba0 · 317aba0
1 parent ae56b9d
commit 317aba0
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
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
@@ -22,6 +22,7 @@
 import com.azure.json.JsonSerializable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -143,19 +144,32 @@ private Mono<?> handleRestResponseReturnType(final HttpResponseDecoder.HttpDecod
                 return response.getSourceResponse().getBody().ignoreElements()
                     .then(Mono.fromCallable(() -> createResponse(response, entityType, null)));
             } else {
-                return handleBodyReturnType(response.getSourceResponse(), response::getDecodedBody, methodParser,
-                    bodyType)
+                return handleBodyReturnType(response.getSourceResponse(),
+                    decodeBytes(response),
+                    methodParser, bodyType)
                     .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject))
                     .switchIfEmpty(Mono.fromCallable(() -> createResponse(response, entityType, null)));
             }
         } else {
             // For now, we're just throwing if the Maybe didn't emit a value.
-            return handleBodyReturnType(response.getSourceResponse(), response::getDecodedBody, methodParser,
-                entityType);
+            return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType);
         }
     }
 
-    static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[], Object> getDecodedBody,
+    private static Function<byte[], Mono<Object>> decodeBytes(HttpResponseDecoder.HttpDecodedResponse response) {
+        return bytes -> Mono.fromCallable(() -> response.getDecodedBody(bytes))
+            .publishOn(Schedulers.boundedElastic())
+            .handle((object, sink) -> {
+                if (object == null) {
+                    sink.complete();
+                } else {
+                    sink.next(object);
+                    sink.complete();
+                }
+            });
+    }
+
+    static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[], Mono<Object>> getDecodedBody,
         SwaggerMethodParser methodParser, Type entityType) {
         final int responseStatusCode = sourceResponse.getStatusCode();
         final HttpMethod httpMethod = methodParser.getHttpMethod();
@@ -197,7 +211,7 @@ static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[]
             asyncResult = sourceResponse.getBodyAsInputStream();
         } else {
             // Mono<Object> or Mono<Page<T>>
-            asyncResult = sourceResponse.getBodyAsByteArray().mapNotNull(getDecodedBody);
+            asyncResult = sourceResponse.getBodyAsByteArray().flatMap(getDecodedBody);
         }
         return asyncResult;
     }

diff --git a/...ore/src/main/resources/META-INF/native-image/com.azure/azure-core/native-image.properties b/...ore/src/main/resources/META-INF/native-image/com.azure/azure-core/native-image.properties
@@ -7,6 +7,7 @@ Args=\
                              ch.qos.logback.classic.Logger,\
                              ch.qos.logback.classic.PatternLayout,\
                              ch.qos.logback.core.CoreConstants,\
+                             ch.qos.logback.core.pattern.parser.Parser,\
                              ch.qos.logback.core.spi.AppenderAttachableImpl,\
                              ch.qos.logback.core.status.InfoStatus,\
                              ch.qos.logback.core.status.StatusBase,\