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());
+ }
+ }
+
}