InternalStreamConnection must release buffers obtained via ByteBufferBsonOutput.getByteBuffers

driver-core/src/main/com/mongodb/connection/Stream.java

@@ -46,7 +46,8 @@ public interface Stream extends BufferProvider{
     /**
      * Write each buffer in the list to the stream in order, blocking until all are completely written.
      *
-     * @param buffers the buffers to write
+     * @param buffers the buffers to write. The operation must not {@linkplain ByteBuf#release() release} any buffer from {@code buffers},
+     * unless it also {@linkplain ByteBuf#retain() retained} it, and releasing is meant to compensate for that.
      * @throws IOException if there are problems writing to the stream
      */
     void write(List<ByteBuf> buffers) throws IOException;
@@ -100,7 +101,8 @@ default ByteBuf read(int numBytes, int additionalTimeout) throws IOException {
      * Write each buffer in the list to the stream in order, asynchronously.  This method should return immediately, and invoke the given
      * callback on completion.
      *
-     * @param buffers the buffers to write
+     * @param buffers the buffers to write. The operation must not {@linkplain ByteBuf#release() release} any buffer from {@code buffers},
+     * unless it also {@linkplain ByteBuf#retain() retained} it, and releasing is meant to compensate for that.
      * @param handler invoked when the write operation has completed
      */
     void writeAsync(List<ByteBuf> buffers, AsyncCompletionHandler<Void> handler);

driver-core/src/main/com/mongodb/connection/netty/NettyStream.java

@@ -249,7 +249,12 @@ public ByteBuf read(final int numBytes, final int additionalTimeoutMillis) throw
     public void writeAsync(final List<ByteBuf> buffers, final AsyncCompletionHandler<Void> handler) {
         CompositeByteBuf composite = PooledByteBufAllocator.DEFAULT.compositeBuffer();
         for (ByteBuf cur : buffers) {
-            composite.addComponent(true, ((NettyByteBuf) cur).asByteBuf());
+            // The Netty framework releases `CompositeByteBuf` after writing
+            // (see https://netty.io/wiki/reference-counted-objects.html#outbound-messages),
+            // which results in the buffer we pass to `CompositeByteBuf.addComponent` being released.
+            // However, `CompositeByteBuf.addComponent` does not retain this buffer,
+            // which means we must retain it to conform to the `Stream.writeAsync` contract.
+            composite.addComponent(true, ((NettyByteBuf) cur).asByteBuf().retain());
         }
 
         channel.writeAndFlush(composite).addListener((ChannelFutureListener) future -> {

driver-core/src/main/com/mongodb/internal/ResourceUtil.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal;
+
+import org.bson.ByteBuf;
+
+/**
+ * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ */
+public final class ResourceUtil {
+    public static void release(final Iterable<? extends ByteBuf> buffers) {
+        // we assume `ByteBuf::release` does not complete abruptly
+        buffers.forEach(buffer -> {
+            if (buffer != null) {
+                buffer.release();
+            }
+        });
+    }
+
+    private ResourceUtil() {
+    }
+}

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java

@@ -42,6 +42,7 @@
 import com.mongodb.connection.Stream;
 import com.mongodb.connection.StreamFactory;
 import com.mongodb.event.CommandListener;
+import com.mongodb.internal.ResourceUtil;
 import com.mongodb.internal.VisibleForTesting;
 import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.internal.diagnostics.logging.Logger;
@@ -412,26 +413,30 @@ private void sendCommandMessage(final CommandMessage message,
 
         Compressor localSendCompressor = sendCompressor;
         if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
+            List<ByteBuf> byteBuffers = bsonOutput.getByteBuffers();
             try {
-                sendMessage(bsonOutput.getByteBuffers(), message.getId());
+                sendMessage(byteBuffers, message.getId());
             } finally {
+                ResourceUtil.release(byteBuffers);
                 bsonOutput.close();
             }
         } else {
-            List<ByteBuf> byteBuffers = bsonOutput.getByteBuffers();
             ByteBufferBsonOutput compressedBsonOutput;
+            List<ByteBuf> byteBuffers = bsonOutput.getByteBuffers();
             try {
                 CompressedMessage compressedMessage = new CompressedMessage(message.getOpCode(), byteBuffers, localSendCompressor,
                         getMessageSettings(description));
                 compressedBsonOutput = new ByteBufferBsonOutput(this);
                 compressedMessage.encode(compressedBsonOutput, sessionContext);
             } finally {
-                releaseAllBuffers(byteBuffers);
+                ResourceUtil.release(byteBuffers);
                 bsonOutput.close();
             }
+            List<ByteBuf> compressedByteBuffers = compressedBsonOutput.getByteBuffers();
             try {
-                sendMessage(compressedBsonOutput.getByteBuffers(), message.getId());
+                sendMessage(compressedByteBuffers, message.getId());
             } finally {
+                ResourceUtil.release(compressedByteBuffers);
                 compressedBsonOutput.close();
             }
         }
@@ -497,7 +502,7 @@ public <T> void sendAndReceiveAsync(final CommandMessage message, final Decoder<
                             getMessageSettings(description));
                     compressedMessage.encode(compressedBsonOutput, sessionContext);
                 } finally {
-                    releaseAllBuffers(byteBuffers);
+                    ResourceUtil.release(byteBuffers);
                     bsonOutput.close();
                 }
                 sendCommandMessageAsync(message.getId(), decoder, sessionContext, callback, compressedBsonOutput, commandEventSender,
@@ -510,16 +515,12 @@ public <T> void sendAndReceiveAsync(final CommandMessage message, final Decoder<
         }
     }
 
-    private void releaseAllBuffers(final List<ByteBuf> byteBuffers) {
-        for (ByteBuf cur : byteBuffers) {
-            cur.release();
-        }
-    }
-
     private <T> void sendCommandMessageAsync(final int messageId, final Decoder<T> decoder, final SessionContext sessionContext,
                                              final SingleResultCallback<T> callback, final ByteBufferBsonOutput bsonOutput,
                                              final CommandEventSender commandEventSender, final boolean responseExpected) {
-        sendMessageAsync(bsonOutput.getByteBuffers(), messageId, (result, t) -> {
+        List<ByteBuf> byteBuffers = bsonOutput.getByteBuffers();
+        sendMessageAsync(byteBuffers, messageId, (result, t) -> {
+            ResourceUtil.release(byteBuffers);
             bsonOutput.close();
             if (t != null) {
                 commandEventSender.sendFailedEvent(t);