Skip to content

Commit d29883f

Browse files
merlimatmerlimat
authored
Refactored the handling of gRPC stubs (#128)
* Refactored the handling of gRPC stubs * Fixed formatting
1 parent b0877ad commit d29883f

21 files changed

+208
-293
lines changed

‎client/src/main/java/io/streamnative/oxia/client/AsyncOxiaClientImpl.java

Lines changed: 13 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -28,24 +28,21 @@
2828
import io.streamnative.oxia.client.batch.Operation.WriteOperation.DeleteOperation;
2929
import io.streamnative.oxia.client.batch.Operation.WriteOperation.DeleteRangeOperation;
3030
import io.streamnative.oxia.client.batch.Operation.WriteOperation.PutOperation;
31-
import io.streamnative.oxia.client.grpc.ChannelManager;
32-
import io.streamnative.oxia.client.grpc.ChannelManager.StubFactory;
31+
import io.streamnative.oxia.client.grpc.OxiaStubManager;
3332
import io.streamnative.oxia.client.metrics.BatchMetrics;
3433
import io.streamnative.oxia.client.metrics.OperationMetrics;
3534
import io.streamnative.oxia.client.notify.NotificationManager;
3635
import io.streamnative.oxia.client.session.SessionManager;
3736
import io.streamnative.oxia.client.shard.ShardManager;
3837
import io.streamnative.oxia.proto.ListRequest;
3938
import io.streamnative.oxia.proto.ListResponse;
40-
import io.streamnative.oxia.proto.ReactorOxiaClientGrpc.ReactorOxiaClientStub;
4139
import java.time.Clock;
4240
import java.util.List;
4341
import java.util.Objects;
4442
import java.util.concurrent.CompletableFuture;
4543
import java.util.concurrent.atomic.AtomicLong;
4644
import java.util.function.Consumer;
4745
import java.util.function.Function;
48-
import java.util.function.Supplier;
4946
import lombok.AccessLevel;
5047
import lombok.NonNull;
5148
import lombok.RequiredArgsConstructor;
@@ -55,16 +52,15 @@
5552
class AsyncOxiaClientImpl implements AsyncOxiaClient {
5653

5754
static @NonNull CompletableFuture<AsyncOxiaClient> newInstance(@NonNull ClientConfig config) {
58-
var channelManager = new ChannelManager();
59-
var reactorStubFactory = channelManager.getReactorStubFactory();
60-
Supplier<ReactorOxiaClientStub> stubFactory =
61-
() -> reactorStubFactory.apply(config.serviceAddress());
62-
var shardManager = new ShardManager(stubFactory, config.metrics(), config.namespace());
63-
var notificationManager =
64-
new NotificationManager(reactorStubFactory, shardManager, config.metrics());
55+
var stubManager = new OxiaStubManager();
56+
57+
var serviceAddrStub = stubManager.getStub(config.serviceAddress());
58+
var shardManager = new ShardManager(serviceAddrStub, config.metrics(), config.namespace());
59+
var notificationManager = new NotificationManager(stubManager, shardManager, config.metrics());
6560

6661
Function<Long, String> leaderFn = shardManager::leader;
67-
var stubByShardId = leaderFn.andThen(reactorStubFactory);
62+
var stubByShardId = leaderFn.andThen(stubManager::getStub);
63+
6864
shardManager.addCallback(notificationManager);
6965
var batchMetrics = BatchMetrics.create(Clock.systemUTC(), config.metrics());
7066
var readBatchManager = BatchManager.newReadBatchManager(config, stubByShardId, batchMetrics);
@@ -76,25 +72,23 @@ class AsyncOxiaClientImpl implements AsyncOxiaClient {
7672

7773
var client =
7874
new AsyncOxiaClientImpl(
79-
channelManager,
75+
stubManager,
8076
shardManager,
8177
notificationManager,
8278
readBatchManager,
8379
writeBatchManager,
8480
sessionManager,
85-
reactorStubFactory,
8681
operationMetrics);
8782

8883
return shardManager.start().thenApply(v -> client);
8984
}
9085

91-
private final @NonNull ChannelManager channelManager;
86+
private final @NonNull OxiaStubManager stubManager;
9287
private final @NonNull ShardManager shardManager;
9388
private final @NonNull NotificationManager notificationManager;
9489
private final @NonNull BatchManager readBatchManager;
9590
private final @NonNull BatchManager writeBatchManager;
9691
private final @NonNull SessionManager sessionManager;
97-
private final @NonNull StubFactory<ReactorOxiaClientStub> reactorStubFactory;
9892
private final @NonNull OperationMetrics metrics;
9993
private final AtomicLong sequence = new AtomicLong();
10094
private volatile boolean closed;
@@ -222,14 +216,14 @@ public void notifications(@NonNull Consumer<Notification> notificationCallback)
222216
private @NonNull Flux<String> list(
223217
long shardId, @NonNull String startKeyInclusive, @NonNull String endKeyExclusive) {
224218
var leader = shardManager.leader(shardId);
225-
var stub = reactorStubFactory.apply(leader);
219+
var stub = stubManager.getStub(leader);
226220
var request =
227221
ListRequest.newBuilder()
228222
.setShardId(shardId)
229223
.setStartInclusive(startKeyInclusive)
230224
.setEndExclusive(endKeyExclusive)
231225
.build();
232-
return stub.list(request).flatMapIterable(ListResponse::getKeysList);
226+
return stub.reactor().list(request).flatMapIterable(ListResponse::getKeysList);
233227
}
234228

235229
@Override
@@ -243,7 +237,7 @@ public void close() throws Exception {
243237
sessionManager.close();
244238
notificationManager.close();
245239
shardManager.close();
246-
channelManager.close();
240+
stubManager.close();
247241
}
248242

249243
private void checkIfClosed() {

‎client/src/main/java/io/streamnative/oxia/client/batch/Batch.java

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,10 @@
2727
import io.streamnative.oxia.client.batch.Operation.WriteOperation.DeleteRangeOperation;
2828
import io.streamnative.oxia.client.batch.Operation.WriteOperation.PutOperation;
2929
import io.streamnative.oxia.client.batch.Operation.WriteOperation.PutOperation.SessionInfo;
30+
import io.streamnative.oxia.client.grpc.OxiaStub;
3031
import io.streamnative.oxia.client.metrics.BatchMetrics;
3132
import io.streamnative.oxia.client.session.SessionManager;
3233
import io.streamnative.oxia.proto.GetResponse;
33-
import io.streamnative.oxia.proto.ReactorOxiaClientGrpc.ReactorOxiaClientStub;
3434
import io.streamnative.oxia.proto.ReadRequest;
3535
import io.streamnative.oxia.proto.WriteRequest;
3636
import java.time.Clock;
@@ -70,7 +70,7 @@ final class WriteBatch extends BatchBase implements Batch {
7070
private long bytes;
7171

7272
WriteBatch(
73-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
73+
@NonNull Function<Long, OxiaStub> stubByShardId,
7474
@NonNull SessionManager sessionManager,
7575
@NonNull String clientIdentifier,
7676
long shardId,
@@ -125,7 +125,7 @@ public void complete() {
125125
sample.startExec();
126126
Throwable t = null;
127127
try {
128-
var response = getStub().write(toProto()).block();
128+
var response = getStub().reactor().write(toProto()).block();
129129
for (var i = 0; i < deletes.size(); i++) {
130130
deletes.get(i).complete(response.getDeletes(i));
131131
}
@@ -180,7 +180,7 @@ public void add(@NonNull Operation<?> operation) {
180180
}
181181

182182
ReadBatch(
183-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
183+
@NonNull Function<Long, OxiaStub> stubByShardId,
184184
long shardId,
185185
long createTime,
186186
BatchMetrics.Sample sample) {
@@ -200,6 +200,7 @@ public void complete() {
200200
try {
201201
var responses =
202202
getStub()
203+
.reactor()
203204
.read(toProto())
204205
.flatMapSequential(response -> Flux.fromIterable(response.getGetsList()))
205206
.doOnNext(r -> bytes.add(r.getValue().size()));
@@ -227,19 +228,19 @@ ReadRequest toProto() {
227228

228229
@RequiredArgsConstructor(access = PRIVATE)
229230
abstract class BatchBase {
230-
private final @NonNull Function<Long, ReactorOxiaClientStub> stubByShardId;
231+
private final @NonNull Function<Long, OxiaStub> stubByShardId;
231232
@Getter private final long shardId;
232233
@Getter private final long startTime;
233234
final BatchMetrics.Sample sample;
234235

235-
protected ReactorOxiaClientStub getStub() {
236+
protected OxiaStub getStub() {
236237
return stubByShardId.apply(shardId);
237238
}
238239
}
239240

240241
@RequiredArgsConstructor(access = PACKAGE)
241242
abstract class BatchFactory implements Function<Long, Batch> {
242-
final @NonNull Function<Long, ReactorOxiaClientStub> stubByShardId;
243+
final @NonNull Function<Long, OxiaStub> stubByShardId;
243244

244245
@Getter(PACKAGE)
245246
private final @NonNull ClientConfig config;
@@ -254,7 +255,7 @@ class WriteBatchFactory extends BatchFactory {
254255
final @NonNull SessionManager sessionManager;
255256

256257
public WriteBatchFactory(
257-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
258+
@NonNull Function<Long, OxiaStub> stubByShardId,
258259
@NonNull SessionManager sessionManager,
259260
@NonNull ClientConfig config,
260261
@NonNull Clock clock,
@@ -278,7 +279,7 @@ public WriteBatchFactory(
278279

279280
class ReadBatchFactory extends BatchFactory {
280281
public ReadBatchFactory(
281-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
282+
@NonNull Function<Long, OxiaStub> stubByShardId,
282283
@NonNull ClientConfig config,
283284
@NonNull Clock clock,
284285
@NonNull BatchMetrics metrics) {

‎client/src/main/java/io/streamnative/oxia/client/batch/BatchManager.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@
2020
import static lombok.AccessLevel.PACKAGE;
2121

2222
import io.streamnative.oxia.client.ClientConfig;
23+
import io.streamnative.oxia.client.grpc.OxiaStub;
2324
import io.streamnative.oxia.client.metrics.BatchMetrics;
2425
import io.streamnative.oxia.client.session.SessionManager;
25-
import io.streamnative.oxia.proto.ReactorOxiaClientGrpc.ReactorOxiaClientStub;
2626
import java.time.Clock;
2727
import java.util.List;
2828
import java.util.Objects;
@@ -99,15 +99,15 @@ public static class ShutdownException extends Exception {
9999

100100
public static @NonNull BatchManager newReadBatchManager(
101101
@NonNull ClientConfig config,
102-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
102+
@NonNull Function<Long, OxiaStub> stubByShardId,
103103
BatchMetrics metrics) {
104104
return new BatchManager(
105105
Batcher.newReadBatcherFactory(config, stubByShardId, Clock.systemUTC(), metrics));
106106
}
107107

108108
public static @NonNull BatchManager newWriteBatchManager(
109109
@NonNull ClientConfig config,
110-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
110+
@NonNull Function<Long, OxiaStub> stubByShardId,
111111
@NonNull SessionManager sessionManager,
112112
BatchMetrics metrics) {
113113
return new BatchManager(

‎client/src/main/java/io/streamnative/oxia/client/batch/Batcher.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@
2020

2121
import io.streamnative.oxia.client.ClientConfig;
2222
import io.streamnative.oxia.client.batch.Operation.CloseOperation;
23+
import io.streamnative.oxia.client.grpc.OxiaStub;
2324
import io.streamnative.oxia.client.metrics.BatchMetrics;
2425
import io.streamnative.oxia.client.session.SessionManager;
25-
import io.streamnative.oxia.proto.ReactorOxiaClientGrpc.ReactorOxiaClientStub;
2626
import java.time.Clock;
2727
import java.util.concurrent.BlockingQueue;
2828
import java.util.concurrent.PriorityBlockingQueue;
@@ -112,7 +112,7 @@ public void run() {
112112

113113
static @NonNull Function<Long, Batcher> newReadBatcherFactory(
114114
@NonNull ClientConfig config,
115-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
115+
@NonNull Function<Long, OxiaStub> stubByShardId,
116116
Clock clock,
117117
BatchMetrics metrics) {
118118
return s ->
@@ -121,7 +121,7 @@ public void run() {
121121

122122
static @NonNull Function<Long, Batcher> newWriteBatcherFactory(
123123
@NonNull ClientConfig config,
124-
@NonNull Function<Long, ReactorOxiaClientStub> stubByShardId,
124+
@NonNull Function<Long, OxiaStub> stubByShardId,
125125
@NonNull SessionManager sessionManager,
126126
Clock clock,
127127
BatchMetrics metrics) {

‎client/src/main/java/io/streamnative/oxia/client/grpc/ChannelManager.java

Lines changed: 0 additions & 96 deletions
This file was deleted.

‎client/src/main/java/io/streamnative/oxia/client/grpc/GrpcResponseStream.java

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,15 @@
1717

1818
import static lombok.AccessLevel.PROTECTED;
1919

20-
import io.streamnative.oxia.proto.ReactorOxiaClientGrpc.ReactorOxiaClientStub;
2120
import java.util.concurrent.CompletableFuture;
2221
import java.util.function.Consumer;
23-
import java.util.function.Supplier;
2422
import lombok.NonNull;
2523
import lombok.RequiredArgsConstructor;
2624
import reactor.core.Disposable;
2725

2826
@RequiredArgsConstructor(access = PROTECTED)
2927
public abstract class GrpcResponseStream implements AutoCloseable {
30-
private final @NonNull Supplier<ReactorOxiaClientStub> stubFactory;
28+
private final @NonNull OxiaStub stub;
3129

3230
private volatile Disposable disposable;
3331

@@ -36,12 +34,12 @@ public abstract class GrpcResponseStream implements AutoCloseable {
3634
if (disposable != null) {
3735
throw new IllegalStateException("Already started");
3836
}
39-
return start(stubFactory.get(), disposable -> this.disposable = disposable);
37+
return start(stub, disposable -> this.disposable = disposable);
4038
}
4139
}
4240

4341
protected abstract @NonNull CompletableFuture<Void> start(
44-
@NonNull ReactorOxiaClientStub stub, @NonNull Consumer<Disposable> consumer);
42+
@NonNull OxiaStub stub, @NonNull Consumer<Disposable> consumer);
4543

4644
@Override
4745
public void close() {

0 commit comments

Comments
 (0)