Replace doFinally with Mono/Flux.using in Core libraries (Azure#36997)

Replace doFinally with Mono/Flux.using in Core libraries

Author: alzimmermsft

eng/versioning/external_dependencies.txt

@@ -114,6 +114,8 @@ redis.clients:jedis;4.3.1
 io.lettuce:lettuce-core;6.2.0.RELEASE
 org.redisson:redisson;3.17.0
 net.bytebuddy:byte-buddy;1.12.23
+testdep_net.bytebuddy:byte-buddy;1.14.8
+testdep_net.bytebuddy:byte-buddy-agent;1.14.8
 
 ## Spring boot dependency versions
 org.springframework.boot:spring-boot-dependencies;2.7.14

sdk/core/azure-core-amqp/pom.xml

@@ -114,6 +114,20 @@
       <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
       <scope>test</scope>
     </dependency>
+    <!-- bytebuddy dependencies are required for mockito 4.11.0 to work with Java 21. Mockito 4.11.0 is the last release -->
+    <!-- of Mockito supporting Java 8 as a baseline. -->
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+      <version>1.14.8</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+      <version>1.14.8</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core-test</artifactId>

sdk/core/azure-core-http-jdk-httpclient/src/main/java/com/azure/core/http/jdk/httpclient/JdkHttpResponseAsync.java

@@ -26,7 +26,7 @@ final class JdkHttpResponseAsync extends JdkHttpResponseBase {
 
     @Override
     public Flux<ByteBuffer> getBody() {
-        return this.contentFlux.doFinally(signalType -> disposed = true);
+        return Flux.using(() -> this, ignored -> contentFlux, ignored -> disposed = true);
     }
 
     @Override

sdk/core/azure-core-http-jdk-httpclient/src/main/java/com/azure/core/http/jdk/httpclient/JdkHttpResponseSync.java

@@ -47,7 +47,7 @@ public Flux<ByteBuffer> getBody() {
         if (bodyBytes != null) {
             return Mono.fromSupplier(() -> ByteBuffer.wrap(bodyBytes)).flux();
         } else {
-            return FluxUtil.toFluxByteBuffer(bodyStream).doFinally(ignored -> close());
+            return Flux.using(() -> this, ignored -> FluxUtil.toFluxByteBuffer(bodyStream), JdkHttpResponseSync::close);
         }
     }
 

sdk/core/azure-core-http-netty/src/main/java/com/azure/core/http/netty/NettyAsyncHttpClient.java

@@ -11,6 +11,7 @@
 import com.azure.core.http.netty.implementation.AzureNettyHttpClientContext;
 import com.azure.core.http.netty.implementation.NettyAsyncHttpBufferedResponse;
 import com.azure.core.http.netty.implementation.NettyAsyncHttpResponse;
+import com.azure.core.http.netty.implementation.Utility;
 import com.azure.core.implementation.util.BinaryDataContent;
 import com.azure.core.implementation.util.BinaryDataHelper;
 import com.azure.core.implementation.util.ByteArrayContent;
@@ -54,8 +55,6 @@
 import java.util.Objects;
 import java.util.function.BiFunction;
 
-import static com.azure.core.http.netty.implementation.Utility.closeConnection;
-
 /**
  * This class provides a Netty-based implementation for the {@link HttpClient} interface. Creating an instance of this
  * class can be achieved by using the {@link NettyAsyncHttpClientBuilder} class, which offers Netty-specific API for
@@ -304,15 +303,14 @@ private static BiFunction<HttpClientResponse, Connection, Mono<Tuple2<HttpRespon
             // For now, eagerlyReadResponse and ignoreResponseBody works the same.
 //            if (ignoreResponseBody) {
 //                AtomicBoolean firstNext = new AtomicBoolean(true);
-//                return reactorNettyConnection.inbound().receive()
+//                return Mono.using(() -> reactorNettyConnection, connection -> connection.inbound().receive()
 //                    .doOnNext(ignored -> {
 //                        if (!firstNext.compareAndSet(true, false)) {
 //                            LOGGER.log(LogLevel.WARNING, () -> "Received HTTP response body when one wasn't expected. "
 //                                + "Response body will be ignored as directed.");
 //                        }
 //                    })
-//                    .ignoreElements()
-//                    .doFinally(ignored -> closeConnection(reactorNettyConnection))
+//                    .ignoreElements(), Utility::closeConnection)
 //                    .then(Mono.fromSupplier(() -> new NettyAsyncHttpBufferedResponse(reactorNettyResponse, restRequest,
 //                        EMPTY_BYTES, headersEagerlyConverted)));
 //            }
@@ -323,11 +321,11 @@ private static BiFunction<HttpClientResponse, Connection, Mono<Tuple2<HttpRespon
              */
             if (eagerlyReadResponse || ignoreResponseBody) {
                 // Set up the body flux and dispose the connection once it has been received.
-                return reactorNettyConnection.inbound().receive().aggregate().asByteArray()
-                    .doFinally(ignored -> closeConnection(reactorNettyConnection))
+                return Mono.using(() -> reactorNettyConnection, connection -> connection.inbound().receive()
+                    .aggregate().asByteArray()
                     .switchIfEmpty(Mono.just(EMPTY_BYTES))
                     .map(bytes -> Tuples.of(new NettyAsyncHttpBufferedResponse(reactorNettyResponse, restRequest, bytes,
-                        headersEagerlyConverted), reactorNettyResponse.responseHeaders()));
+                        headersEagerlyConverted), reactorNettyResponse.responseHeaders())), Utility::closeConnection);
             } else {
                 return Mono.just(Tuples.of(new NettyAsyncHttpResponse(reactorNettyResponse, reactorNettyConnection,
                     restRequest, disableBufferCopy, headersEagerlyConverted), reactorNettyResponse.responseHeaders()));

sdk/core/azure-core-http-netty/src/main/java/com/azure/core/http/netty/implementation/NettyAsyncHttpResponse.java

@@ -38,14 +38,15 @@ public NettyAsyncHttpResponse(HttpClientResponse reactorNettyResponse, Connectio
 
     @Override
     public Flux<ByteBuffer> getBody() {
-        return bodyIntern()
-            .map(byteBuf -> this.disableBufferCopy ? byteBuf.nioBuffer() : deepCopyBuffer(byteBuf))
-            .doFinally(ignored -> close());
+        return Flux.using(() -> this, response -> response.bodyIntern()
+            .map(byteBuf -> this.disableBufferCopy ? byteBuf.nioBuffer() : deepCopyBuffer(byteBuf)),
+            NettyAsyncHttpResponse::close);
     }
 
     @Override
     public Mono<byte[]> getBodyAsByteArray() {
-        return bodyIntern().aggregate().asByteArray().doFinally(ignored -> close());
+        return Mono.using(() -> this, response -> response.bodyIntern().aggregate().asByteArray(),
+            NettyAsyncHttpResponse::close);
     }
 
     @Override
@@ -56,12 +57,14 @@ public Mono<String> getBodyAsString() {
 
     @Override
     public Mono<String> getBodyAsString(Charset charset) {
-        return bodyIntern().aggregate().asString(charset).doFinally(ignored -> close());
+        return Mono.using(() -> this, response -> response.bodyIntern().aggregate().asString(charset),
+            NettyAsyncHttpResponse::close);
     }
 
     @Override
     public Mono<InputStream> getBodyAsInputStream() {
-        return bodyIntern().aggregate().asInputStream().doFinally(ignored -> close());
+        return Mono.using(() -> this, response -> response.bodyIntern().aggregate().asInputStream(),
+            NettyAsyncHttpResponse::close);
     }
 
     @Override
@@ -89,10 +92,9 @@ public void writeBodyTo(WritableByteChannel channel) {
         // complete. This introduces a previously seen, but in a different flavor, race condition where the write
         // operation gets scheduled on one thread and the ByteBuf release happens on another, leaving the write
         // operation racing to complete before the release happens. With all that said, leave this as subscribeOn.
-        Mono.<Void>create(sink -> bodyIntern().subscribe(
+        Mono.using(() -> this, response -> Mono.<Void>create(sink -> response.bodyIntern().subscribe(
             new ByteBufWriteSubscriber(channel::write, sink, getContentLength())))
-            .subscribeOn(Schedulers.boundedElastic())
-            .doFinally(ignored -> close())
+            .subscribeOn(Schedulers.boundedElastic()), NettyAsyncHttpResponse::close)
             .block();
     }
 

sdk/core/azure-core-http-vertx/pom.xml

@@ -141,6 +141,20 @@
       <version>4.11.0</version><!-- {x-version-update;org.mockito:mockito-inline;external_dependency} -->
       <scope>test</scope>
     </dependency>
+    <!-- bytebuddy dependencies are required for mockito 4.11.0 to work with Java 21. Mockito 4.11.0 is the last release -->
+    <!-- of Mockito supporting Java 8 as a baseline. -->
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+      <version>1.14.8</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+      <version>1.14.8</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

sdk/core/azure-core-test/src/main/java/com/azure/core/test/http/HttpClientTests.java

@@ -2366,16 +2366,14 @@ public void simpleDownloadTest(Context context) {
     @MethodSource("downloadTestArgumentProvider")
     public void simpleDownloadTestAsync(Context context) {
         StepVerifier.create(createService(DownloadService.class).getBytesAsync(getRequestUri(), context)
-                .flatMap(response -> response.getValue().map(ByteBuffer::remaining)
-                    .reduce(0, Integer::sum)
-                    .doFinally(ignore -> response.close())))
+                .flatMap(response -> Mono.using(() -> response, r -> r.getValue().map(ByteBuffer::remaining)
+                    .reduce(0, Integer::sum), StreamResponse::close)))
             .assertNext(count -> assertEquals(30720, count))
             .verifyComplete();
 
         StepVerifier.create(createService(DownloadService.class).getBytesAsync(getRequestUri(), context)
-                .flatMap(response -> Mono.zip(MessageDigestUtils.md5(response.getValue()),
-                        Mono.just(response.getHeaders().getValue(HttpHeaderName.ETAG)))
-                    .doFinally(ignore -> response.close())))
+                .flatMap(response -> Mono.using(() -> response, r -> Mono.zip(MessageDigestUtils.md5(r.getValue()),
+                    Mono.just(r.getHeaders().getValue(HttpHeaderName.ETAG))), StreamResponse::close)))
             .assertNext(hashTuple -> assertEquals(hashTuple.getT2(), hashTuple.getT1()))
             .verifyComplete();
     }
@@ -2433,12 +2431,13 @@ public void streamResponseCanTransferBodyAsync(Context context) throws IOExcepti
                         () -> IOUtils.toAsynchronousByteChannel(AsynchronousFileChannel.open(tempFile, StandardOpenOption.WRITE), 0),
                         streamResponse::writeValueToAsync,
                         channel -> {
+                            streamResponse.close();
                             try {
                                 channel.close();
                             } catch (IOException e) {
                                 throw Exceptions.propagate(e);
                             }
-                        }).doFinally(ignored -> streamResponse.close())
+                        })
                     .then(Mono.just(streamResponse.getHeaders().getValue(HttpHeaderName.ETAG)))))
             .assertNext(hash -> {
                 try {

sdk/core/azure-core/src/main/java/com/azure/core/credential/SimpleTokenCache.java

@@ -85,8 +85,8 @@ public Mono<AccessToken> getToken() {
                         // cache hasn't expired, ignore refresh error this time
                         fallback = Mono.just(cache);
                     }
-                    return tokenRefresh
-                        .materialize()
+
+                    return Mono.using(() -> wip, ignored -> tokenRefresh.materialize()
                         .flatMap(signal -> {
                             AccessToken accessToken = signal.get();
                             Throwable error = signal.getThrowable();
@@ -107,8 +107,7 @@ public Mono<AccessToken> getToken() {
                                 return fallback;
                             }
                         })
-                        .doOnError(sinksOne::tryEmitError)
-                        .doFinally(ignored -> wip.set(null));
+                        .doOnError(sinksOne::tryEmitError), w -> w.set(null));
                 } else if (cache != null && !cache.isExpired()) {
                     // another thread might be refreshing the token proactively, but the current token is still valid
                     return Mono.just(cache);

sdk/core/azure-core/src/main/java/com/azure/core/http/HttpResponse.java

@@ -175,8 +175,7 @@ public Mono<Void> writeBodyToAsync(AsynchronousByteChannel channel) {
     public void writeBodyTo(WritableByteChannel channel) throws IOException {
         Flux<ByteBuffer> body = getBody();
         if (body != null) {
-            FluxUtil.writeToWritableByteChannel(body, channel)
-                .doFinally(ignored -> close())
+            Mono.using(() -> this, ignored -> FluxUtil.writeToWritableByteChannel(body, channel), HttpResponse::close)
                 .block();
         }
     }