fix: reference counts in PerChannelBookieClient

Author: lhotari

bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/PerChannelBookieClient.java

@@ -777,6 +777,8 @@ void addEntry(final long ledgerId, byte[] masterKey, final long entryId, Referen
                   Object ctx, final int options, boolean allowFastFail, final EnumSet<WriteFlag> writeFlags) {
         Object request = null;
         CompletionKey completionKey = null;
+        Runnable cleanupActionBeforeWrite = null;
+        Runnable cleanupActionAfterWrite = null;
         if (useV2WireProtocol) {
             if (writeFlags.contains(WriteFlag.DEFERRED_SYNC)) {
                 LOG.error("invalid writeflags {} for v2 protocol", writeFlags);
@@ -786,9 +788,13 @@ void addEntry(final long ledgerId, byte[] masterKey, final long entryId, Referen
             completionKey = acquireV2Key(ledgerId, entryId, OperationType.ADD_ENTRY);
 
             if (toSend instanceof ByteBuf) {
-                request = ((ByteBuf) toSend).retainedDuplicate();
+                ByteBuf byteBuf = ((ByteBuf) toSend).retainedDuplicate();
+                request = byteBuf;
+                cleanupActionBeforeWrite = byteBuf::release;
             } else {
-                request = ByteBufList.clone((ByteBufList) toSend);
+                ByteBufList byteBufList = ByteBufList.clone((ByteBufList) toSend);
+                request = byteBufList;
+                cleanupActionBeforeWrite = byteBufList::release;
             }
         } else {
             final long txnId = getTxnId();
@@ -805,7 +811,9 @@ void addEntry(final long ledgerId, byte[] masterKey, final long entryId, Referen
 
             ByteString body = null;
             ByteBufList bufToSend = (ByteBufList) toSend;
-
+            bufToSend.retain();
+            cleanupActionBeforeWrite = bufToSend::release;
+            cleanupActionAfterWrite = cleanupActionBeforeWrite;
             if (bufToSend.hasArray()) {
                 body = UnsafeByteOperations.unsafeWrap(bufToSend.array(), bufToSend.arrayOffset(),
                         bufToSend.readableBytes());
@@ -840,17 +848,9 @@ void addEntry(final long ledgerId, byte[] masterKey, final long entryId, Referen
         putCompletionKeyValue(completionKey,
                               acquireAddCompletion(completionKey,
                                                    cb, ctx, ledgerId, entryId));
-        final Channel c = channel;
-        if (c == null) {
-            // usually checked in writeAndFlush, but we have extra check
-            // because we need to release toSend.
-            errorOut(completionKey);
-            ReferenceCountUtil.release(toSend);
-            return;
-        } else {
-            // addEntry times out on backpressure
-            writeAndFlush(c, completionKey, request, allowFastFail);
-        }
+        // addEntry times out on backpressure
+        writeAndFlush(channel, completionKey, request, allowFastFail, cleanupActionBeforeWrite,
+                cleanupActionAfterWrite);
     }
 
     public void readLac(final long ledgerId, ReadLacCallback cb, Object ctx) {
@@ -1005,7 +1005,7 @@ private void readEntryInternal(final long ledgerId,
         ReadCompletion readCompletion = new ReadCompletion(completionKey, cb, ctx, ledgerId, entryId);
         putCompletionKeyValue(completionKey, readCompletion);
 
-        writeAndFlush(channel, completionKey, request, allowFastFail);
+        writeAndFlush(channel, completionKey, request, allowFastFail, null, null);
     }
 
     public void batchReadEntries(final long ledgerId,
@@ -1048,7 +1048,7 @@ private void batchReadEntriesInternal(final long ledgerId,
                 completionKey, cb, ctx, ledgerId, startEntryId);
         putCompletionKeyValue(completionKey, readCompletion);
 
-        writeAndFlush(channel, completionKey, request, allowFastFail);
+        writeAndFlush(channel, completionKey, request, allowFastFail, null, null);
     }
 
     public void getBookieInfo(final long requested, GetBookieInfoCallback cb, Object ctx) {
@@ -1170,16 +1170,20 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exceptio
     private void writeAndFlush(final Channel channel,
                                final CompletionKey key,
                                final Object request) {
-        writeAndFlush(channel, key, request, false);
+        writeAndFlush(channel, key, request, false, null, null);
     }
 
     private void writeAndFlush(final Channel channel,
                            final CompletionKey key,
                            final Object request,
-                           final boolean allowFastFail) {
+                               final boolean allowFastFail, final Runnable cleanupActionBeforeWrite,
+                               final Runnable cleanupActionAfterWrite) {
         if (channel == null) {
             LOG.warn("Operation {} failed: channel == null", StringUtils.requestToString(request));
             errorOut(key);
+            if (cleanupActionBeforeWrite != null) {
+                cleanupActionBeforeWrite.run();
+            }
             return;
         }
 
@@ -1194,27 +1198,39 @@ private void writeAndFlush(final Channel channel,
                     StringUtils.requestToString(request));
 
             errorOut(key, BKException.Code.TooManyRequestsException);
+            if (cleanupActionBeforeWrite != null) {
+                cleanupActionBeforeWrite.run();
+            }
             return;
         }
 
         try {
             final long startTime = MathUtils.nowInNano();
 
             ChannelPromise promise = channel.newPromise().addListener(future -> {
-                if (future.isSuccess()) {
-                    nettyOpLogger.registerSuccessfulEvent(MathUtils.elapsedNanos(startTime), TimeUnit.NANOSECONDS);
-                    CompletionValue completion = completionObjects.get(key);
-                    if (completion != null) {
-                        completion.setOutstanding();
+                try {
+                    if (future.isSuccess()) {
+                        nettyOpLogger.registerSuccessfulEvent(MathUtils.elapsedNanos(startTime), TimeUnit.NANOSECONDS);
+                        CompletionValue completion = completionObjects.get(key);
+                        if (completion != null) {
+                            completion.setOutstanding();
+                        }
+                    } else {
+                        nettyOpLogger.registerFailedEvent(MathUtils.elapsedNanos(startTime), TimeUnit.NANOSECONDS);
+                    }
+                } finally {
+                    if (cleanupActionAfterWrite != null) {
+                        cleanupActionAfterWrite.run();
                     }
-                } else {
-                    nettyOpLogger.registerFailedEvent(MathUtils.elapsedNanos(startTime), TimeUnit.NANOSECONDS);
                 }
             });
             channel.writeAndFlush(request, promise);
         } catch (Throwable e) {
             LOG.warn("Operation {} failed", StringUtils.requestToString(request), e);
             errorOut(key);
+            if (cleanupActionBeforeWrite != null) {
+                cleanupActionBeforeWrite.run();
+            }
         }
     }