ByteBuffer is not released when compression is enabled, causing pool leakage (#883)

JAVA-4510

Co-authored-by: Konstantin Labun <[email protected]>
Co-authored-by: Valentin Kovalenko <[email protected]>

Author: 3 people

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

@@ -724,19 +724,27 @@ private ResponseBuffers receiveResponseBuffers(final int additionalTimeout) thro
         }
 
         ByteBuf messageBuffer = stream.read(messageHeader.getMessageLength() - MESSAGE_HEADER_LENGTH, additionalTimeout);
+        boolean releaseMessageBuffer = true;
+        try {
+            if (messageHeader.getOpCode() == OP_COMPRESSED.getValue()) {
+                CompressedHeader compressedHeader = new CompressedHeader(messageBuffer, messageHeader);
 
-        if (messageHeader.getOpCode() == OP_COMPRESSED.getValue()) {
-            CompressedHeader compressedHeader = new CompressedHeader(messageBuffer, messageHeader);
-
-            Compressor compressor = getCompressor(compressedHeader);
+                Compressor compressor = getCompressor(compressedHeader);
 
-            ByteBuf buffer = getBuffer(compressedHeader.getUncompressedSize());
-            compressor.uncompress(messageBuffer, buffer);
+                ByteBuf buffer = getBuffer(compressedHeader.getUncompressedSize());
+                compressor.uncompress(messageBuffer, buffer);
 
-            buffer.flip();
-            return new ResponseBuffers(new ReplyHeader(buffer, compressedHeader), buffer);
-        } else {
-            return new ResponseBuffers(new ReplyHeader(messageBuffer, messageHeader), messageBuffer);
+                buffer.flip();
+                return new ResponseBuffers(new ReplyHeader(buffer, compressedHeader), buffer);
+            } else {
+                ResponseBuffers responseBuffers = new ResponseBuffers(new ReplyHeader(messageBuffer, messageHeader), messageBuffer);
+                releaseMessageBuffer = false;
+                return responseBuffers;
+            }
+        } finally {
+            if (releaseMessageBuffer) {
+                messageBuffer.release();
+            }
         }
     }
 
@@ -792,6 +800,7 @@ public void onResult(final ByteBuf result, final Throwable t) {
                     callback.onResult(null, t);
                     return;
                 }
+                boolean releaseResult = true;
                 try {
                     ReplyHeader replyHeader;
                     ByteBuf responseBuffer;
@@ -806,15 +815,21 @@ public void onResult(final ByteBuf result, final Throwable t) {
                             replyHeader = new ReplyHeader(buffer, compressedHeader);
                             responseBuffer = buffer;
                         } finally {
+                            releaseResult = false;
                             result.release();
                         }
                     } else {
                         replyHeader = new ReplyHeader(result, messageHeader);
                         responseBuffer = result;
+                        releaseResult = false;
                     }
                     callback.onResult(new ResponseBuffers(replyHeader, responseBuffer), null);
                 } catch (Throwable localThrowable) {
                     callback.onResult(null, localThrowable);
+                } finally {
+                    if (releaseResult) {
+                        result.release();
+                    }
                 }
             }
         }