Skip to content

perf: serialize write requests outside the write stream lock #435

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merlimat merged 1 commit into from
Oct 6, 2026
+144 −21
Merged

perf: serialize write requests outside the write stream lock #435

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 28 additions & 5 deletions client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -66,6 +68,25 @@ final class GrpcRpcProvider implements RpcProvider {
private static final Metadata.Key<String> 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<InputStream, WriteResponse> WRITE_STREAM_METHOD =
OxiaClientGrpc.getWriteStreamMethod().toBuilder(
new MethodDescriptor.Marshaller<InputStream>() {
@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;
Expand Down Expand Up @@ -335,7 +356,7 @@ public ManagedWriteStream getWriteStream(long shardId) {
}

@Override
public StreamObserver<WriteRequest> writeStream(
public StreamObserver<InputStream> writeStream(
long shardId,
OxiaStatusException leaderHint,
StreamObserver<WriteResponse> responseObserver) {
Expand All @@ -345,12 +366,14 @@ public StreamObserver<WriteRequest> 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();
Expand Down
14 changes: 11 additions & 3 deletions client/src/main/java/io/oxia/client/grpc/ManagedSubWriteStream.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<WriteResponse> {

private final ManagedWriteStream parent;
private final StreamObserver<WriteRequest> requestObserver;
private final StreamObserver<InputStream> requestObserver;

ManagedSubWriteStream(
ManagedWriteStream parent,
Expand All @@ -33,8 +35,14 @@ final class ManagedSubWriteStream implements StreamObserver<WriteResponse> {
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() {
Expand Down
9 changes: 6 additions & 3 deletions client/src/main/java/io/oxia/client/grpc/ManagedWriteStream.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.*;
Expand Down Expand Up @@ -196,9 +197,11 @@ void handleCompleted(ManagedSubWriteStream source) {
}

public CompletableFuture<WriteResponse> 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<WriteResponse> future = new CompletableFuture<>();
final InflightWrite inflightWrite = new InflightWrite(request, future, System.nanoTime());
lock.lock();
try {
log.debug(
event ->
Expand All @@ -223,7 +226,7 @@ public CompletableFuture<WriteResponse> send(WriteRequest request) {
initWithRecovery(null);
}
} else {
subStreamObserver.send(inflightWrite.request);
subStreamObserver.send(serializedRequest);
log.debug(
event ->
event
Expand Down Expand Up @@ -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 ->
Expand Down
4 changes: 3 additions & 1 deletion client/src/main/java/io/oxia/client/grpc/RpcProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -63,7 +64,8 @@ CompletableFuture<KeepAliveResponse> keepAlive(

void read(@NonNull ReadRequest request, @NonNull StreamObserver<ReadResponse> observer);

StreamObserver<WriteRequest> writeStream(
/** Opens a write stream, which takes requests serialized by ManagedSubWriteStream.serialize(). */
StreamObserver<InputStream> writeStream(
long shardId, OxiaStatusException leaderHint, StreamObserver<WriteResponse> responseObserver);

ManagedWriteStream getWriteStream(long shardId);
Expand Down
105 changes: 96 additions & 9 deletions client/src/test/java/io/oxia/client/grpc/ManagedWriteStreamTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -78,6 +84,66 @@ void completesWritesInResponseOrder() throws Exception {
}
}

@Test
void handlesResponseWhileAnotherWriteIsSerialized() throws Exception {
var requests = new LinkedBlockingQueue<WriteRequest>();
var responses = new CompletableFuture<StreamObserver<WriteResponse>>();
Server server =
writeServer(
new OxiaClientGrpc.OxiaClientImplBase() {
@Override
public StreamObserver<WriteRequest> writeStream(
StreamObserver<WriteResponse> 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 =
Expand Down Expand Up @@ -131,9 +197,9 @@ void ignoresSynchronousResponseDuringClose() throws Exception {
.thenAnswer(
invocation -> {
var responseObserver = invocation.<StreamObserver<WriteResponse>>getArgument(2);
return new StreamObserver<WriteRequest>() {
return new StreamObserver<InputStream>() {
@Override
public void onNext(WriteRequest value) {}
public void onNext(InputStream value) {}

@Override
public void onError(Throwable t) {}
Expand Down Expand Up @@ -292,10 +358,10 @@ void resetsObserverWhenRecoveryReplayFails() throws Exception {
var responseObserver = invocation.<StreamObserver<WriteResponse>>getArgument(2);
responseObservers.add(responseObserver);
var attempt = openAttempts.incrementAndGet();
return new StreamObserver<WriteRequest>() {
return new StreamObserver<InputStream>() {
@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();
Expand Down Expand Up @@ -351,9 +417,9 @@ void ignoresResponsesFromInactiveSubStream() throws Exception {
invocation -> {
var responseObserver = invocation.<StreamObserver<WriteResponse>>getArgument(2);
responseObservers.add(responseObserver);
return new StreamObserver<WriteRequest>() {
return new StreamObserver<InputStream>() {
@Override
public void onNext(WriteRequest value) {
public void onNext(InputStream value) {
requestCount.incrementAndGet();
}

Expand Down Expand Up @@ -438,9 +504,9 @@ void backsOffReconnectsUntilStreamGetsResponse() throws Exception {
.thenAnswer(
invocation -> {
responseObservers.add(invocation.getArgument(2));
return new StreamObserver<WriteRequest>() {
return new StreamObserver<InputStream>() {
@Override
public void onNext(WriteRequest value) {}
public void onNext(InputStream value) {}

@Override
public void onError(Throwable t) {}
Expand Down Expand Up @@ -546,6 +612,27 @@ private static java.util.List<String> values(Queue<WriteRequest> 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)
Expand Down
Loading