This is an automated email from the ASF dual-hosted git repository.

merlimat pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/bookkeeper.git


The following commit(s) were added to refs/heads/master by this push:
     new 451b6a59ef Fix table service txn request corruption on silent client 
retries (#4829)
451b6a59ef is described below

commit 451b6a59ef93352f791002ffaa7159a53e5399c2
Author: Lari Hotari <[email protected]>
AuthorDate: Fri Jul 10 20:57:43 2026 +0300

    Fix table service txn request corruption on silent client retries (#4829)
    
    TableClientTest.testTableAPIServerSideRouting has been flaky since 2018
    (#1440), failing at assertTrue(txnResult.isSuccess()) with no error in
    the client logs.
    
    Root cause: table service requests embed ByteBuf slices and serializing
    a request drains them, so a request instance must not be reused across
    RPC attempts. Both txn client paths reused the same TxnRequest across
    silent retries: PByteBufSimpleTableImpl.TxnImpl.commit() passed one
    pre-populated request into RetryUtils.execute (server-side routing) and
    TxnRequestProcessor.createRequest() returned the same stored request on
    every attempt (client-side routing). A transient first-attempt failure
    (e.g. a cold storage-container proxy channel returning NOT_FOUND by
    design) made the retry serialize the drained request, sending corrupted
    bytes. The server then answered gRPC-OK with code=BAD_REQUEST and
    succeeded=false, which KvUtils.newKvTxnResult conflated with a
    legitimately failed compare because it never checks the response code.
    put/get/delete/increment already build a fresh request per attempt;
    txn was the outlier.
    
    Changes:
    - Store the compare/success/failure ops in both TxnImpl classes and
      build a fresh TxnRequest for every RPC attempt.
    - TxnRequestProcessor now takes a Supplier<TxnRequest>.
    - PByteBufSimpleTableImpl surfaces non-SUCCESS txn response codes as
      InternalServerException instead of isSuccess()=false, matching
      TxnRequestProcessor.processResponse.
    - Regression tests that fail against the previous implementation.
    
    Assisted-by: Claude Code
---
 .../clients/impl/kv/PByteBufSimpleTableImpl.java   |  62 ++++++--
 .../clients/impl/kv/PByteBufTableRangeImpl.java    |  50 +++++--
 .../clients/impl/kv/TxnRequestProcessor.java       |  15 +-
 .../impl/kv/PByteBufSimpleTableImplTest.java       | 166 +++++++++++++++++++++
 .../clients/impl/kv/TxnRequestProcessorTest.java   |  81 +++++++++-
 5 files changed, 338 insertions(+), 36 deletions(-)

diff --git 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImpl.java
 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImpl.java
index ae699c1f39..12f2d0f8ef 100644
--- 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImpl.java
+++ 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImpl.java
@@ -64,7 +64,9 @@ import org.apache.bookkeeper.api.kv.result.IncrementResult;
 import org.apache.bookkeeper.api.kv.result.PutResult;
 import org.apache.bookkeeper.api.kv.result.RangeResult;
 import org.apache.bookkeeper.api.kv.result.TxnResult;
+import org.apache.bookkeeper.clients.exceptions.InternalServerException;
 import org.apache.bookkeeper.clients.utils.RetryUtils;
+import org.apache.bookkeeper.common.concurrent.FutureUtils;
 import org.apache.bookkeeper.stream.proto.StreamProperties;
 import org.apache.bookkeeper.stream.proto.kv.rpc.DeleteRangeRequest;
 import org.apache.bookkeeper.stream.proto.kv.rpc.IncrementRequest;
@@ -72,6 +74,7 @@ import org.apache.bookkeeper.stream.proto.kv.rpc.PutRequest;
 import org.apache.bookkeeper.stream.proto.kv.rpc.RangeRequest;
 import org.apache.bookkeeper.stream.proto.kv.rpc.RoutingHeader;
 import org.apache.bookkeeper.stream.proto.kv.rpc.TxnRequest;
+import org.apache.bookkeeper.stream.proto.storage.StatusCode;
 
 /**
  * A {@link PTable} implementation using simple grpc calls.
@@ -260,21 +263,22 @@ public class PByteBufSimpleTableImpl
     class TxnImpl implements Txn<ByteBuf, ByteBuf> {
 
         private final ByteBuf pKey;
-        private final TxnRequest txnRequest;
-        private final List<AutoCloseable> resourcesToRelease;
+        private final List<CompareOp<ByteBuf, ByteBuf>> compareOps;
+        private final List<Op<ByteBuf, ByteBuf>> successOps;
+        private final List<Op<ByteBuf, ByteBuf>> failureOps;
 
         TxnImpl(ByteBuf pKey) {
             this.pKey = pKey.retain();
-            this.txnRequest = new TxnRequest();
-            this.resourcesToRelease = Lists.newArrayList();
+            this.compareOps = Lists.newArrayList();
+            this.successOps = Lists.newArrayList();
+            this.failureOps = Lists.newArrayList();
         }
 
         @SuppressWarnings("unchecked")
         @Override
         public Txn<ByteBuf, ByteBuf> If(CompareOp... cmps) {
             for (CompareOp<ByteBuf, ByteBuf> cmp : cmps) {
-                populateProtoCompare(txnRequest.addCompare(), cmp);
-                resourcesToRelease.add(cmp);
+                compareOps.add(cmp);
             }
             return this;
         }
@@ -283,8 +287,7 @@ public class PByteBufSimpleTableImpl
         @Override
         public Txn<ByteBuf, ByteBuf> Then(Op... ops) {
             for (Op<ByteBuf, ByteBuf> op : ops) {
-                populateProtoRequest(txnRequest.addSuccess(), op);
-                resourcesToRelease.add(op);
+                successOps.add(op);
             }
             return this;
         }
@@ -293,25 +296,54 @@ public class PByteBufSimpleTableImpl
         @Override
         public Txn<ByteBuf, ByteBuf> Else(Op... ops) {
             for (Op<ByteBuf, ByteBuf> op : ops) {
-                populateProtoRequest(txnRequest.addFailure(), op);
-                resourcesToRelease.add(op);
+                failureOps.add(op);
             }
             return this;
         }
 
+        // Serializing a request drains the ByteBuf slices stored in it, so a 
request instance
+        // must not be reused across RPC attempts: a retried attempt would 
send a corrupted
+        // request. Build a fresh request per attempt, like 
put/get/delete/increment above.
+        private TxnRequest newTxnRequest() {
+            TxnRequest txnRequest = new TxnRequest();
+            for (CompareOp<ByteBuf, ByteBuf> cmp : compareOps) {
+                populateProtoCompare(txnRequest.addCompare(), cmp);
+            }
+            for (Op<ByteBuf, ByteBuf> op : successOps) {
+                populateProtoRequest(txnRequest.addSuccess(), op);
+            }
+            for (Op<ByteBuf, ByteBuf> op : failureOps) {
+                populateProtoRequest(txnRequest.addFailure(), op);
+            }
+            populateRoutingHeader(txnRequest.setHeader(), pKey);
+            return txnRequest;
+        }
+
         @Override
         public CompletableFuture<TxnResult<ByteBuf, ByteBuf>> commit() {
-            populateRoutingHeader(txnRequest.setHeader(), pKey);
             return retryUtils.execute(() -> fromListenableFuture(
             ClientCalls.futureUnaryCall(
                 getChannel(pKey).newCall(getTxnMethod(), getCallOptions()),
-                txnRequest)
+                newTxnRequest())
             ))
-            .thenApply(response -> KvUtils.newKvTxnResult(response, 
resultFactory, kvFactory))
+            .thenCompose(response -> {
+                if (StatusCode.SUCCESS != response.getHeader().getCode()) {
+                    // A server-side error must not be conflated with a failed 
txn compare.
+                    return FutureUtils.exception(new InternalServerException(
+                        "Encountered internal server exception : code = " + 
response.getHeader().getCode()));
+                }
+                return FutureUtils.value(KvUtils.newKvTxnResult(response, 
resultFactory, kvFactory));
+            })
             .whenComplete((ignored, cause) -> {
                 ReferenceCountUtil.release(pKey);
-                for (AutoCloseable resource : resourcesToRelease) {
-                    closeResource(resource);
+                for (CompareOp<ByteBuf, ByteBuf> cmp : compareOps) {
+                    closeResource(cmp);
+                }
+                for (Op<ByteBuf, ByteBuf> op : successOps) {
+                    closeResource(op);
+                }
+                for (Op<ByteBuf, ByteBuf> op : failureOps) {
+                    closeResource(op);
                 }
             });
         }
diff --git 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufTableRangeImpl.java
 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufTableRangeImpl.java
index 2875d90675..6a72045b83 100644
--- 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufTableRangeImpl.java
+++ 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/PByteBufTableRangeImpl.java
@@ -209,21 +209,22 @@ class PByteBufTableRangeImpl implements PTable<ByteBuf, 
ByteBuf> {
     class TxnImpl implements Txn<ByteBuf, ByteBuf> {
 
         private final ByteBuf pKey;
-        private final TxnRequest txnRequest;
-        private final List<AutoCloseable> resourcesToRelease;
+        private final List<CompareOp<ByteBuf, ByteBuf>> compareOps;
+        private final List<Op<ByteBuf, ByteBuf>> successOps;
+        private final List<Op<ByteBuf, ByteBuf>> failureOps;
 
         TxnImpl(ByteBuf pKey) {
             this.pKey = pKey.retain();
-            this.txnRequest = new TxnRequest();
-            this.resourcesToRelease = Lists.newArrayList();
+            this.compareOps = Lists.newArrayList();
+            this.successOps = Lists.newArrayList();
+            this.failureOps = Lists.newArrayList();
         }
 
         @SuppressWarnings("unchecked")
         @Override
         public Txn<ByteBuf, ByteBuf> If(CompareOp... cmps) {
             for (CompareOp<ByteBuf, ByteBuf> cmp : cmps) {
-                populateProtoCompare(txnRequest.addCompare(), cmp);
-                resourcesToRelease.add(cmp);
+                compareOps.add(cmp);
             }
             return this;
         }
@@ -232,8 +233,7 @@ class PByteBufTableRangeImpl implements PTable<ByteBuf, 
ByteBuf> {
         @Override
         public Txn<ByteBuf, ByteBuf> Then(Op... ops) {
             for (Op<ByteBuf, ByteBuf> op : ops) {
-                populateProtoRequest(txnRequest.addSuccess(), op);
-                resourcesToRelease.add(op);
+                successOps.add(op);
             }
             return this;
         }
@@ -242,25 +242,47 @@ class PByteBufTableRangeImpl implements PTable<ByteBuf, 
ByteBuf> {
         @Override
         public Txn<ByteBuf, ByteBuf> Else(Op... ops) {
             for (Op<ByteBuf, ByteBuf> op : ops) {
-                populateProtoRequest(txnRequest.addFailure(), op);
-                resourcesToRelease.add(op);
+                failureOps.add(op);
             }
             return this;
         }
 
+        // Serializing a request drains the ByteBuf slices stored in it, so a 
request instance
+        // must not be reused across RPC attempts: a retried attempt would 
send a corrupted
+        // request. Build a fresh request per attempt.
+        private TxnRequest newTxnRequest() {
+            TxnRequest txnRequest = new TxnRequest();
+            for (CompareOp<ByteBuf, ByteBuf> cmp : compareOps) {
+                populateProtoCompare(txnRequest.addCompare(), cmp);
+            }
+            for (Op<ByteBuf, ByteBuf> op : successOps) {
+                populateProtoRequest(txnRequest.addSuccess(), op);
+            }
+            for (Op<ByteBuf, ByteBuf> op : failureOps) {
+                populateProtoRequest(txnRequest.addFailure(), op);
+            }
+            populateRoutingHeader(txnRequest.setHeader(), pKey);
+            return txnRequest;
+        }
+
         @Override
         public CompletableFuture<TxnResult<ByteBuf, ByteBuf>> commit() {
-            populateRoutingHeader(txnRequest.setHeader(), pKey);
             return TxnRequestProcessor.of(
-                txnRequest,
+                this::newTxnRequest,
                 response -> KvUtils.newKvTxnResult(response, resultFactory, 
kvFactory),
                 scChannel,
                 executor,
                 backoffPolicy
             ).process().whenComplete((ignored, cause) -> {
                 ReferenceCountUtil.release(pKey);
-                for (AutoCloseable resource : resourcesToRelease) {
-                    closeResource(resource);
+                for (CompareOp<ByteBuf, ByteBuf> cmp : compareOps) {
+                    closeResource(cmp);
+                }
+                for (Op<ByteBuf, ByteBuf> op : successOps) {
+                    closeResource(op);
+                }
+                for (Op<ByteBuf, ByteBuf> op : failureOps) {
+                    closeResource(op);
                 }
             });
         }
diff --git 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessor.java
 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessor.java
index a3ed20925c..72363f09c1 100644
--- 
a/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessor.java
+++ 
b/stream/clients/java/kv/src/main/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessor.java
@@ -35,6 +35,7 @@ package org.apache.bookkeeper.clients.impl.kv;
 import com.google.common.util.concurrent.ListenableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import org.apache.bookkeeper.clients.exceptions.InternalServerException;
 import org.apache.bookkeeper.clients.impl.channel.StorageServerChannel;
 import org.apache.bookkeeper.clients.impl.container.StorageContainerChannel;
@@ -51,30 +52,32 @@ class TxnRequestProcessor<RespT>
     extends ListenableFutureRpcProcessor<TxnRequest, TxnResponse, RespT> {
 
     public static <T> TxnRequestProcessor<T> of(
-        TxnRequest request,
+        Supplier<TxnRequest> requestSupplier,
         Function<TxnResponse, T> responseFunc,
         StorageContainerChannel channel,
         ScheduledExecutorService executor,
         Policy backoffPolicy) {
-        return new TxnRequestProcessor<>(request, responseFunc, channel, 
executor, backoffPolicy);
+        return new TxnRequestProcessor<>(requestSupplier, responseFunc, 
channel, executor, backoffPolicy);
     }
 
-    private final TxnRequest request;
+    private final Supplier<TxnRequest> requestSupplier;
     private final Function<TxnResponse, RespT> responseFunc;
 
-    private TxnRequestProcessor(TxnRequest request,
+    private TxnRequestProcessor(Supplier<TxnRequest> requestSupplier,
                                 Function<TxnResponse, RespT> respFunc,
                                 StorageContainerChannel channel,
                                 ScheduledExecutorService executor,
                                 Policy backoffPolicy) {
         super(channel, executor, backoffPolicy);
-        this.request = request;
+        this.requestSupplier = requestSupplier;
         this.responseFunc = respFunc;
     }
 
     @Override
     protected TxnRequest createRequest() {
-        return request;
+        // Serializing a request drains the ByteBuf slices stored in it, so a 
request instance
+        // must not be reused across RPC attempts: build a fresh request for 
every attempt.
+        return requestSupplier.get();
     }
 
     @Override
diff --git 
a/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImplTest.java
 
b/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImplTest.java
new file mode 100644
index 0000000000..02e9d8dfeb
--- /dev/null
+++ 
b/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/PByteBufSimpleTableImplTest.java
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.bookkeeper.clients.impl.kv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import io.grpc.CallOptions;
+import io.grpc.Status;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.stub.StreamObserver;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
+import org.apache.bookkeeper.api.kv.PTable;
+import org.apache.bookkeeper.api.kv.Txn;
+import org.apache.bookkeeper.api.kv.op.CompareResult;
+import org.apache.bookkeeper.api.kv.result.TxnResult;
+import org.apache.bookkeeper.clients.exceptions.InternalServerException;
+import org.apache.bookkeeper.clients.grpc.GrpcClientTestBase;
+import org.apache.bookkeeper.clients.utils.ClientConstants;
+import org.apache.bookkeeper.clients.utils.RetryUtils;
+import org.apache.bookkeeper.common.concurrent.FutureUtils;
+import org.apache.bookkeeper.stream.proto.StreamProperties;
+import 
org.apache.bookkeeper.stream.proto.kv.rpc.TableServiceGrpc.TableServiceImplBase;
+import org.apache.bookkeeper.stream.proto.kv.rpc.TxnRequest;
+import org.apache.bookkeeper.stream.proto.kv.rpc.TxnResponse;
+import org.apache.bookkeeper.stream.proto.storage.StatusCode;
+import org.junit.Test;
+
+/**
+ * Unit test of {@link PByteBufSimpleTableImpl}.
+ */
+public class PByteBufSimpleTableImplTest extends GrpcClientTestBase {
+
+    private static final long STREAM_ID = 1234L;
+
+    private StreamProperties streamProps;
+
+    @Override
+    protected void doSetup() throws Exception {
+        streamProps = new StreamProperties();
+        streamProps.setStreamId(STREAM_ID);
+    }
+
+    @Override
+    protected void doTeardown() throws Exception {
+    }
+
+    private PTable<ByteBuf, ByteBuf> newTable() {
+        return new PByteBufSimpleTableImpl(
+            streamProps,
+            
InProcessChannelBuilder.forName(serverName).directExecutor().build(),
+            CallOptions.DEFAULT,
+            RetryUtils.create(ClientConstants.DEFAULT_BACKOFF_POLICY, 
scheduler));
+    }
+
+    /**
+     * Serializing a request drains the ByteBuf slices stored in it, so a 
request instance must
+     * not be reused across RPC attempts. Verify that a retried txn commit 
sends an intact
+     * request, rebuilt for every attempt.
+     */
+    @Test
+    public void testTxnRetrySendsIntactRequest() throws Exception {
+        List<TxnRequest> receivedRequests = new CopyOnWriteArrayList<>();
+        TableServiceImplBase tableService = new TableServiceImplBase() {
+            @Override
+            public void txn(TxnRequest request,
+                            StreamObserver<TxnResponse> responseObserver) {
+                receivedRequests.add(request);
+                if (receivedRequests.size() == 1) {
+                    // fail the first attempt with a retryable status to force 
a retry
+                    
responseObserver.onError(Status.UNAVAILABLE.asRuntimeException());
+                } else {
+                    TxnResponse response = new TxnResponse();
+                    response.setHeader().setCode(StatusCode.SUCCESS);
+                    response.setSucceeded(true);
+                    responseObserver.onNext(response);
+                    responseObserver.onCompleted();
+                }
+            }
+        };
+        serviceRegistry.addService(tableService.bindService());
+
+        PTable<ByteBuf, ByteBuf> table = newTable();
+        ByteBuf key = Unpooled.wrappedBuffer("txn-key".getBytes(UTF_8));
+        ByteBuf value = Unpooled.wrappedBuffer("txn-value".getBytes(UTF_8));
+        Txn<ByteBuf, ByteBuf> txn = table.txn(key);
+        CompletableFuture<TxnResult<ByteBuf, ByteBuf>> commitFuture = txn
+            .If(
+                table.opFactory().compareValue(CompareResult.EQUAL, key, 
Unpooled.wrappedBuffer(new byte[0]))
+            )
+            .Then(
+                table.opFactory().newPut(key, value, 
table.opFactory().optionFactory().newPutOption().build())
+            )
+            .commit();
+        try (TxnResult<ByteBuf, ByteBuf> txnResult = 
FutureUtils.result(commitFuture)) {
+            assertTrue(txnResult.isSuccess());
+        }
+
+        assertEquals(2, receivedRequests.size());
+        for (TxnRequest received : receivedRequests) {
+            assertEquals(1, received.getComparesCount());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getCompareAt(0).getKey());
+            assertEquals(1, received.getSuccessesCount());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getSuccessAt(0).getRequestPut().getKey());
+            assertArrayEquals("txn-value".getBytes(UTF_8), 
received.getSuccessAt(0).getRequestPut().getValue());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getHeader().getRKey());
+            assertEquals(STREAM_ID, received.getHeader().getStreamId());
+        }
+    }
+
+    /**
+     * A server-side error response must surface as an exception instead of 
being conflated
+     * with a legitimately failed txn compare (isSuccess() == false).
+     */
+    @Test
+    public void testTxnServerErrorNotConflatedWithFailedCompare() throws 
Exception {
+        TableServiceImplBase tableService = new TableServiceImplBase() {
+            @Override
+            public void txn(TxnRequest request,
+                            StreamObserver<TxnResponse> responseObserver) {
+                TxnResponse response = new TxnResponse();
+                response.setHeader().setCode(StatusCode.INTERNAL_SERVER_ERROR);
+                responseObserver.onNext(response);
+                responseObserver.onCompleted();
+            }
+        };
+        serviceRegistry.addService(tableService.bindService());
+
+        PTable<ByteBuf, ByteBuf> table = newTable();
+        ByteBuf key = Unpooled.wrappedBuffer("txn-key".getBytes(UTF_8));
+        CompletableFuture<TxnResult<ByteBuf, ByteBuf>> commitFuture = 
table.txn(key)
+            .If(
+                table.opFactory().compareValue(CompareResult.EQUAL, key, 
Unpooled.wrappedBuffer(new byte[0]))
+            )
+            .commit();
+        try {
+            FutureUtils.result(commitFuture);
+            fail("txn commit should fail when the server responds with an 
error code");
+        } catch (InternalServerException e) {
+            // expected
+        }
+    }
+
+}
diff --git 
a/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessorTest.java
 
b/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessorTest.java
index 6da33ad4d1..d36c67cb34 100644
--- 
a/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessorTest.java
+++ 
b/stream/clients/java/kv/src/test/java/org/apache/bookkeeper/clients/impl/kv/TxnRequestProcessorTest.java
@@ -17,17 +17,25 @@
  */
 package org.apache.bookkeeper.clients.impl.kv;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import io.grpc.Status;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.stub.StreamObserver;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import lombok.Cleanup;
 import org.apache.bookkeeper.clients.grpc.GrpcClientTestBase;
 import org.apache.bookkeeper.clients.impl.channel.StorageServerChannel;
@@ -98,7 +106,7 @@ public class TxnRequestProcessorTest extends 
GrpcClientTestBase {
         TxnRequest request = newRequest();
 
         TxnRequestProcessor<String> processor = TxnRequestProcessor.of(
-            request,
+            () -> request,
             resp -> "test",
             scChannel,
             scheduler,
@@ -107,4 +115,75 @@ public class TxnRequestProcessorTest extends 
GrpcClientTestBase {
         assertEquals(request, receivedRequest.get());
     }
 
+    /**
+     * Serializing a request drains the ByteBuf slices stored in it, so a 
request instance must
+     * not be reused across RPC attempts. Verify that a retried txn sends an 
intact request built
+     * fresh from the request supplier for every attempt.
+     */
+    @Test
+    public void testRequestRebuiltForEachAttempt() throws Exception {
+        @Cleanup("shutdown") ScheduledExecutorService scheduler = 
Executors.newSingleThreadScheduledExecutor();
+
+        StorageContainerChannel scChannel = 
mock(StorageContainerChannel.class);
+
+        CompletableFuture<StorageServerChannel> serverChannelFuture = 
FutureUtils.createFuture();
+        
when(scChannel.getStorageContainerChannelFuture()).thenReturn(serverChannelFuture);
+
+        TxnResponse response = newSuccessResponse();
+
+        List<TxnRequest> receivedRequests = new CopyOnWriteArrayList<>();
+        TableServiceImplBase tableService = new TableServiceImplBase() {
+
+            @Override
+            public void txn(TxnRequest request,
+                            StreamObserver<TxnResponse> responseObserver) {
+                receivedRequests.add(request);
+                if (receivedRequests.size() == 1) {
+                    // fail the first attempt with a retryable status to force 
a retry
+                    
responseObserver.onError(Status.NOT_FOUND.asRuntimeException());
+                } else {
+                    responseObserver.onNext(response);
+                    responseObserver.onCompleted();
+                }
+            }
+        };
+        serviceRegistry.addService(tableService.bindService());
+        StorageServerChannel ssChannel = new StorageServerChannel(
+            
InProcessChannelBuilder.forName(serverName).directExecutor().build(),
+            Optional.empty());
+        serverChannelFuture.complete(ssChannel);
+
+        ByteBuf key = Unpooled.wrappedBuffer("txn-key".getBytes(UTF_8));
+        ByteBuf value = Unpooled.wrappedBuffer("txn-value".getBytes(UTF_8));
+        Supplier<TxnRequest> requestSupplier = () -> {
+            TxnRequest request = new TxnRequest();
+            request.addCompare().setKey(key.slice());
+            request.addSuccess().setRequestPut()
+                .setKey(key.slice())
+                .setValue(value.slice());
+            request.setHeader()
+                .setStreamId(1234L)
+                .setRKey(key.slice());
+            return request;
+        };
+
+        TxnRequestProcessor<String> processor = TxnRequestProcessor.of(
+            requestSupplier,
+            resp -> "test",
+            scChannel,
+            scheduler,
+            ClientConstants.DEFAULT_INFINIT_BACKOFF_POLICY);
+        assertEquals("test", FutureUtils.result(processor.process()));
+
+        assertEquals(2, receivedRequests.size());
+        for (TxnRequest received : receivedRequests) {
+            assertEquals(1, received.getComparesCount());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getCompareAt(0).getKey());
+            assertEquals(1, received.getSuccessesCount());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getSuccessAt(0).getRequestPut().getKey());
+            assertArrayEquals("txn-value".getBytes(UTF_8), 
received.getSuccessAt(0).getRequestPut().getValue());
+            assertArrayEquals("txn-key".getBytes(UTF_8), 
received.getHeader().getRKey());
+        }
+    }
+
 }

Reply via email to