Fixes ByteBuf Leak during Edge Cases of writeBodyTo (Azure#30670)

Fixes ByteBuf Leak during Edge Cases of writeBodyTo

Author: alzimmermsft

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

@@ -40,8 +40,9 @@ public NettyAsyncHttpResponse(HttpClientResponse reactorNettyResponse, Connectio
 
     @Override
     public Flux<ByteBuffer> getBody() {
-        return bodyIntern().doFinally(ignored -> close())
-            .map(byteBuf -> this.disableBufferCopy ? byteBuf.nioBuffer() : deepCopyBuffer(byteBuf));
+        return bodyIntern()
+            .map(byteBuf -> this.disableBufferCopy ? byteBuf.nioBuffer() : deepCopyBuffer(byteBuf))
+            .doFinally(ignored -> close());
     }
 
     @Override
@@ -61,7 +62,7 @@ public Mono<String> getBodyAsString(Charset charset) {
 
     @Override
     public Mono<InputStream> getBodyAsInputStream() {
-        return bodyIntern().aggregate().asInputStream();
+        return bodyIntern().aggregate().asInputStream().doFinally(ignored -> close());
     }
 
     @Override
@@ -70,13 +71,28 @@ public Mono<Void> writeBodyToAsync(AsynchronousByteChannel channel) {
             .flatMapSequential(nettyBuffer ->
                 FluxUtil.writeToAsynchronousByteChannel(Flux.just(nettyBuffer.nioBuffer()), channel)
                     .doFinally(ignored -> nettyBuffer.release()), 1, 1)
+            .doFinally(ignored -> close())
             .then();
     }
 
     @Override
     public void writeBodyTo(WritableByteChannel channel) {
-        bodyIntern().retain()
-            .publishOn(Schedulers.boundedElastic())
+        // Since this uses a synchronous write this doesn't need to retain the ByteBuf as the async version does.
+        // In fact, if retain is used here and there is a cancellation or exception during writing this will likely
+        // leak ByteBufs as it doesn't have the asynchronous stream handling to properly close them when errors happen
+        // during writing.
+        //
+        // This also must use subscribeOn rather than publishOn as subscribeOn will have the entire asynchronous chain
+        // run on the same subscriber where publishOn only affects reactive operations after the publishOn. While this
+        // doesn't seem like it would make much of a difference as this entire call will be blocked it fixes a race
+        // condition. reactor-netty's inbound receive uses a ByteBuf pool to handle the network response and the
+        // returned ByteBufFlux is configured to release the ByteBuf back to the pool when the onNext operation
+        // completes. Unfortunately, when publishOn is used each ByteBuf must be scheduled in the publisher thread using
+        // a future and the publishOn onNext handler will complete once that operation is scheduled, not when it's
+        // 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.
+        bodyIntern().subscribeOn(Schedulers.boundedElastic())
             .map(nettyBuffer -> {
                 try {
                     ByteBuffer nioBuffer = nettyBuffer.nioBuffer();
@@ -86,10 +102,10 @@ public void writeBodyTo(WritableByteChannel channel) {
                     return nettyBuffer;
                 } catch (IOException e) {
                     throw Exceptions.propagate(e);
-                } finally {
-                    nettyBuffer.release();
                 }
-            }).then().block();
+            })
+            .doFinally(ignored -> close())
+            .then().block();
     }
 
     @Override

sdk/core/azure-core-http-netty/src/test/java/com/azure/core/http/netty/implementation/HttpResponseDrainsBufferTests.java

@@ -24,10 +24,17 @@
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.AsynchronousByteChannel;
+import java.nio.channels.CompletionHandler;
+import java.nio.channels.WritableByteChannel;
 import java.security.SecureRandom;
 import java.time.Duration;
 import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 
@@ -106,6 +113,99 @@ public void closeHttpResponseWithConsumingPartialBody() {
         runScenario(response -> response.getBody().next().flatMap(ignored -> Mono.fromRunnable(response::close)));
     }
 
+    @Test
+    public void closeHttpResponseWithConsumingPartialWriteAsync() {
+        runScenario(response -> response.writeBodyToAsync(new ThrowingAsynchronousByteChannel())
+            .onErrorResume(throwable -> Mono.empty()));
+    }
+
+    private static final class ThrowingAsynchronousByteChannel implements AsynchronousByteChannel {
+        private boolean open = true;
+        int writeCount = 0;
+
+        @Override
+        public <A> void read(ByteBuffer dst, A attachment, CompletionHandler<Integer, ? super A> handler) {
+        }
+
+        @Override
+        public Future<Integer> read(ByteBuffer dst) {
+            return null;
+        }
+
+        @Override
+        public <A> void write(ByteBuffer src, A attachment, CompletionHandler<Integer, ? super A> handler) {
+            if (writeCount++ < 3) {
+                int remaining = src.remaining();
+                src.position(src.position() + remaining);
+                handler.completed(remaining, attachment);
+            } else {
+                handler.failed(new IOException(), attachment);
+            }
+        }
+
+        @Override
+        public Future<Integer> write(ByteBuffer src) {
+            if (writeCount++ < 3) {
+                int remaining = src.remaining();
+                src.position(src.position() + remaining);
+
+                return CompletableFuture.completedFuture(remaining);
+            } else {
+                CompletableFuture<Integer> failed = new CompletableFuture<>();
+                failed.completeExceptionally(new IOException());
+                return failed;
+            }
+        }
+
+        @Override
+        public boolean isOpen() {
+            return open;
+        }
+
+        @Override
+        public void close() {
+            open = false;
+        }
+    }
+
+    @Test
+    public void closeHttpResponseWithConsumingPartialWrite() {
+        runScenario(response -> {
+            try {
+                response.writeBodyTo(new ThrowingWritableByteChannel());
+            } catch (Exception ignored) {
+            }
+
+            return Mono.empty();
+        });
+    }
+
+    private static final class ThrowingWritableByteChannel implements WritableByteChannel {
+        private boolean open = true;
+        int writeCount = 0;
+
+        @Override
+        public int write(ByteBuffer src) throws IOException {
+            if (writeCount++ < 3) {
+                int remaining = src.remaining();
+                src.position(src.position() + remaining);
+                return remaining;
+            } else {
+                throw new IOException();
+            }
+        }
+
+        @Override
+        public boolean isOpen() {
+            return open;
+        }
+
+        @Override
+        public void close() throws IOException {
+            open = false;
+        }
+    }
+
     @Test
     public void closeHttpResponseWithConsumingFullBody() {
         runScenario(response -> response.getBodyAsByteArray().flatMap(ignored -> Mono.fromRunnable(response::close)));

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

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