diff --git a/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java b/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java index 489badbb..50f9cbf0 100644 --- a/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java +++ b/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java @@ -22,6 +22,8 @@ import io.github.merlimat.slog.Logger; import io.grpc.Context; import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.stub.ClientCalls; import io.grpc.stub.MetadataUtils; import io.grpc.stub.StreamObserver; import io.oxia.client.ClientConfig; @@ -47,8 +49,8 @@ import io.oxia.proto.SessionHeartbeat; import io.oxia.proto.ShardAssignments; import io.oxia.proto.ShardAssignmentsRequest; -import io.oxia.proto.WriteRequest; import io.oxia.proto.WriteResponse; +import java.io.InputStream; import java.time.Duration; import java.util.Map; import java.util.concurrent.CompletableFuture; @@ -66,6 +68,25 @@ final class GrpcRpcProvider implements RpcProvider { private static final Metadata.Key SHARD_ID_KEY = Metadata.Key.of("shard-id", Metadata.ASCII_STRING_MARSHALLER); + // The write stream takes requests that ManagedSubWriteStream.serialize() serialized, and sends + // them as they are, so that serializing them doesn't hold the ManagedWriteStream lock. Each one + // can only be streamed once, which holds as the channel doesn't retry calls. + private static final MethodDescriptor WRITE_STREAM_METHOD = + OxiaClientGrpc.getWriteStreamMethod().toBuilder( + new MethodDescriptor.Marshaller() { + @Override + public InputStream stream(InputStream value) { + return value; + } + + @Override + public InputStream parse(InputStream stream) { + return stream; + } + }, + OxiaClientGrpc.getWriteStreamMethod().getResponseMarshaller()) + .build(); + private final ClientConfig clientConfig; private final ConnectionManager connectionManager; private final boolean ownsConnectionManager; @@ -335,7 +356,7 @@ public ManagedWriteStream getWriteStream(long shardId) { } @Override - public StreamObserver writeStream( + public StreamObserver writeStream( long shardId, OxiaStatusException leaderHint, StreamObserver responseObserver) { @@ -345,12 +366,14 @@ public StreamObserver writeStream( final var headers = new Metadata(); headers.put(NAMESPACE_KEY, clientConfig.namespace()); headers.put(SHARD_ID_KEY, Long.toString(shardId)); - final var requestObserver = + final var stub = connectionManager .getConnection(getLeader(shardId, hint)) .stub() - .withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers)) - .writeStream(barrierObserver); + .withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers)); + final var requestObserver = + ClientCalls.asyncBidiStreamingCall( + stub.getChannel().newCall(WRITE_STREAM_METHOD, stub.getCallOptions()), barrierObserver); // we don't need to wait for the stream to be ready, only need to check if it's fast failed future.complete(null); future.join(); diff --git a/client/src/main/java/io/oxia/client/grpc/ManagedSubWriteStream.java b/client/src/main/java/io/oxia/client/grpc/ManagedSubWriteStream.java index b00403b3..3cc95725 100644 --- a/client/src/main/java/io/oxia/client/grpc/ManagedSubWriteStream.java +++ b/client/src/main/java/io/oxia/client/grpc/ManagedSubWriteStream.java @@ -16,13 +16,15 @@ package io.oxia.client.grpc; import io.grpc.stub.StreamObserver; +import io.oxia.proto.OxiaClientGrpc; import io.oxia.proto.WriteRequest; import io.oxia.proto.WriteResponse; +import java.io.InputStream; final class ManagedSubWriteStream implements StreamObserver { private final ManagedWriteStream parent; - private final StreamObserver requestObserver; + private final StreamObserver requestObserver; ManagedSubWriteStream( ManagedWriteStream parent, @@ -33,8 +35,14 @@ final class ManagedSubWriteStream implements StreamObserver { this.requestObserver = rpcProvider.writeStream(shardId, leaderHint, this); } - void send(WriteRequest request) { - requestObserver.onNext(request); + // Serializes a request for send(), so that the caller can do it before taking its lock. The + // result can be sent once: sending it reads it + static InputStream serialize(WriteRequest request) { + return OxiaClientGrpc.getWriteStreamMethod().streamRequest(request); + } + + void send(InputStream serializedRequest) { + requestObserver.onNext(serializedRequest); } void complete() { diff --git a/client/src/main/java/io/oxia/client/grpc/ManagedWriteStream.java b/client/src/main/java/io/oxia/client/grpc/ManagedWriteStream.java index c83edded..63816a9e 100644 --- a/client/src/main/java/io/oxia/client/grpc/ManagedWriteStream.java +++ b/client/src/main/java/io/oxia/client/grpc/ManagedWriteStream.java @@ -19,6 +19,7 @@ import io.oxia.client.util.Backoff; import io.oxia.proto.WriteRequest; import io.oxia.proto.WriteResponse; +import java.io.InputStream; import java.time.Duration; import java.util.*; import java.util.concurrent.*; @@ -196,9 +197,11 @@ void handleCompleted(ManagedSubWriteStream source) { } public CompletableFuture send(WriteRequest request) { - lock.lock(); + // Serialize the request before taking the lock, which every response for the shard needs + final InputStream serializedRequest = ManagedSubWriteStream.serialize(request); final CompletableFuture future = new CompletableFuture<>(); final InflightWrite inflightWrite = new InflightWrite(request, future, System.nanoTime()); + lock.lock(); try { log.debug( event -> @@ -223,7 +226,7 @@ public CompletableFuture send(WriteRequest request) { initWithRecovery(null); } } else { - subStreamObserver.send(inflightWrite.request); + subStreamObserver.send(serializedRequest); log.debug( event -> event @@ -316,7 +319,7 @@ private void initWithRecovery(OxiaStatusException leaderHint) { .attr("leaderHint", leaderHint) .log("Replaying inflight writes on opened stream")); for (InflightWrite inflightWrite : inflightWrites) { - subStreamObserver.send(inflightWrite.request); + subStreamObserver.send(ManagedSubWriteStream.serialize(inflightWrite.request)); } log.debug( event -> diff --git a/client/src/main/java/io/oxia/client/grpc/RpcProvider.java b/client/src/main/java/io/oxia/client/grpc/RpcProvider.java index b367af21..85fbb9eb 100644 --- a/client/src/main/java/io/oxia/client/grpc/RpcProvider.java +++ b/client/src/main/java/io/oxia/client/grpc/RpcProvider.java @@ -19,6 +19,7 @@ import io.oxia.client.ClientConfig; import io.oxia.client.grpc.observer.CancelableStreamObserver; import io.oxia.proto.*; +import java.io.InputStream; import java.time.Duration; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; @@ -63,7 +64,8 @@ CompletableFuture keepAlive( void read(@NonNull ReadRequest request, @NonNull StreamObserver observer); - StreamObserver writeStream( + /** Opens a write stream, which takes requests serialized by ManagedSubWriteStream.serialize(). */ + StreamObserver writeStream( long shardId, OxiaStatusException leaderHint, StreamObserver responseObserver); ManagedWriteStream getWriteStream(long shardId); diff --git a/client/src/test/java/io/oxia/client/grpc/ManagedWriteStreamTest.java b/client/src/test/java/io/oxia/client/grpc/ManagedWriteStreamTest.java index efddd3fd..0899866b 100644 --- a/client/src/test/java/io/oxia/client/grpc/ManagedWriteStreamTest.java +++ b/client/src/test/java/io/oxia/client/grpc/ManagedWriteStreamTest.java @@ -32,15 +32,21 @@ import io.grpc.Status; import io.grpc.protobuf.StatusProto; import io.grpc.stub.StreamObserver; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.UnpooledByteBufAllocator; +import io.netty.buffer.UnpooledHeapByteBuf; import io.oxia.client.OxiaClientBuilderImpl; import io.oxia.client.api.OxiaClientBuilder; import io.oxia.proto.OxiaClientGrpc; import io.oxia.proto.WriteRequest; import io.oxia.proto.WriteResponse; +import java.io.InputStream; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.Queue; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledFuture; @@ -78,6 +84,66 @@ void completesWritesInResponseOrder() throws Exception { } } + @Test + void handlesResponseWhileAnotherWriteIsSerialized() throws Exception { + var requests = new LinkedBlockingQueue(); + var responses = new CompletableFuture>(); + Server server = + writeServer( + new OxiaClientGrpc.OxiaClientImplBase() { + @Override + public StreamObserver writeStream( + StreamObserver responseObserver) { + responses.complete(responseObserver); + return new StreamObserver<>() { + @Override + public void onNext(WriteRequest value) { + requests.add(value); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + } + }); + var address = "localhost:" + server.getPort(); + var executor = Executors.newSingleThreadScheduledExecutor(); + var sender = Executors.newSingleThreadExecutor(); + var config = clientConfig(address); + var value = new BlockingValue("value-2".getBytes(StandardCharsets.UTF_8)); + + try (var provider = new GrpcRpcProvider(config, executor, shard -> address); + var stream = new ManagedWriteStream(1, provider, executor, config.requestTimeout())) { + var first = stream.send(writeRequest(1)); + assertThat(requests.poll(5, TimeUnit.SECONDS)).isNotNull(); + + var request = new WriteRequest().setShard(1); + request.addPut().setKey("key-2").setValue(value); + var second = sender.submit(() -> stream.send(request)); + try { + assertThat(value.serializing.await(5, TimeUnit.SECONDS)).isTrue(); + // The response to the first write doesn't wait for the second one to be serialized + responses.get().onNext(new WriteResponse()); + first.get(5, TimeUnit.SECONDS); + } finally { + value.proceed.countDown(); + } + + var secondRequest = requests.poll(5, TimeUnit.SECONDS); + assertThat(secondRequest.getPutAt(0).getValue()) + .isEqualTo("value-2".getBytes(StandardCharsets.UTF_8)); + responses.get().onNext(new WriteResponse()); + second.get(5, TimeUnit.SECONDS).get(5, TimeUnit.SECONDS); + } finally { + sender.shutdownNow(); + executor.shutdownNow(); + server.shutdownNow(); + } + } + @Test void closeFailsPendingWritesAndRejectsNewWrites() throws Exception { Server server = @@ -131,9 +197,9 @@ void ignoresSynchronousResponseDuringClose() throws Exception { .thenAnswer( invocation -> { var responseObserver = invocation.>getArgument(2); - return new StreamObserver() { + return new StreamObserver() { @Override - public void onNext(WriteRequest value) {} + public void onNext(InputStream value) {} @Override public void onError(Throwable t) {} @@ -292,10 +358,10 @@ void resetsObserverWhenRecoveryReplayFails() throws Exception { var responseObserver = invocation.>getArgument(2); responseObservers.add(responseObserver); var attempt = openAttempts.incrementAndGet(); - return new StreamObserver() { + return new StreamObserver() { @Override - public void onNext(WriteRequest value) { - requests.add(value); + public void onNext(InputStream value) { + requests.add(OxiaClientGrpc.getWriteStreamMethod().parseRequest(value)); if (attempt == 2) { replayFailures.incrementAndGet(); throw Status.UNAVAILABLE.withDescription("failed replay").asRuntimeException(); @@ -351,9 +417,9 @@ void ignoresResponsesFromInactiveSubStream() throws Exception { invocation -> { var responseObserver = invocation.>getArgument(2); responseObservers.add(responseObserver); - return new StreamObserver() { + return new StreamObserver() { @Override - public void onNext(WriteRequest value) { + public void onNext(InputStream value) { requestCount.incrementAndGet(); } @@ -438,9 +504,9 @@ void backsOffReconnectsUntilStreamGetsResponse() throws Exception { .thenAnswer( invocation -> { responseObservers.add(invocation.getArgument(2)); - return new StreamObserver() { + return new StreamObserver() { @Override - public void onNext(WriteRequest value) {} + public void onNext(InputStream value) {} @Override public void onError(Throwable t) {} @@ -546,6 +612,27 @@ private static java.util.List values(Queue requests) { .toList(); } + // A value whose serialization waits until the test lets it proceed + private static final class BlockingValue extends UnpooledHeapByteBuf { + final CountDownLatch serializing = new CountDownLatch(1); + final CountDownLatch proceed = new CountDownLatch(1); + + BlockingValue(byte[] value) { + super(UnpooledByteBufAllocator.DEFAULT, value, value.length); + } + + @Override + public ByteBuf getBytes(int index, byte[] dst, int dstIndex, int length) { + serializing.countDown(); + try { + proceed.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return super.getBytes(index, dst, dstIndex, length); + } + } + private static io.oxia.client.ClientConfig clientConfig(String address) { return ((OxiaClientBuilderImpl) OxiaClientBuilder.create(address)