Skip to content

Commit 2c250d4

Browse files
mutianfmutianf
authored
refactor(bigtable): make sessionMap shared by permissions (#14662)
1 parent 2a314aa commit 2c250d4

11 files changed

Lines changed: 462 additions & 212 deletions

File tree

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadModifyWriteRowShim;
4141
import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim;
4242
import com.google.cloud.bigtable.data.v2.internal.compat.ops.RowBuilderShim;
43+
import com.google.cloud.bigtable.data.v2.internal.compat.ops.SessionPoolMap;
4344
import com.google.cloud.bigtable.data.v2.internal.csm.Metrics;
4445
import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo;
4546
import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer;
@@ -85,6 +86,7 @@ public class ShimImpl implements Shim {
8586
private final Resource<ClientConfigurationManager> configManagerResource;
8687
private final Client client;
8788
private final DebugTagTracer debugTagTracer;
89+
private final SessionPoolMap sessionMap;
8890

8991
private final ReadRowShim readRowShim;
9092
private final MutateRowShim mutateRowShim;
@@ -202,11 +204,12 @@ public ShimImpl(
202204
this.configManager = configManagerResource.get();
203205
this.client = client;
204206
this.debugTagTracer = debugTagTracer;
207+
this.sessionMap = new SessionPoolMap(client);
205208

206-
this.readRowShim = new ReadRowShim(client);
207-
this.mutateRowShim = new MutateRowShim(client);
208-
this.checkAndMutateRowShim = new CheckAndMutateRowShim(client);
209-
this.readModifyWriteRowShim = new ReadModifyWriteRowShim(client);
209+
this.readRowShim = new ReadRowShim(sessionMap);
210+
this.mutateRowShim = new MutateRowShim(sessionMap);
211+
this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap);
212+
this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap);
210213
}
211214

212215
/**
@@ -358,6 +361,7 @@ private static boolean checkDirectAccessAvailable(
358361

359362
@Override
360363
public void close() {
364+
sessionMap.invalidateAll();
361365
client.close();
362366
configManagerResource.close();
363367
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import io.grpc.Context;
2626
import io.grpc.Deadline;
2727
import java.io.Closeable;
28-
import java.io.IOException;
2928
import java.time.Duration;
3029
import java.util.Optional;
3130
import java.util.concurrent.TimeUnit;
@@ -46,7 +45,7 @@ public static <K, V extends Closeable> LoadingCache<K, V> createSessionMap(
4645
V value = notification.getValue();
4746
//noinspection DataFlowIssue
4847
value.close();
49-
} catch (IOException e) {
48+
} catch (Exception e) {
5049
// TODO log
5150
}
5251
})

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java‎

Lines changed: 13 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -15,58 +15,37 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.internal.compat.ops;
1717

18-
import com.google.bigtable.v2.OpenAuthorizedViewRequest;
19-
import com.google.bigtable.v2.OpenTableRequest.Permission;
2018
import com.google.bigtable.v2.SessionCheckAndMutateRowRequest;
21-
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
22-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
23-
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2419
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2520
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2621
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
2722
import com.google.cloud.bigtable.data.v2.models.TableId;
2823
import com.google.cloud.bigtable.data.v2.models.TargetId;
2924
import io.grpc.Deadline;
30-
import java.io.IOException;
3125
import java.util.concurrent.CompletableFuture;
3226

3327
public class CheckAndMutateRowShim implements UnaryShim<ConditionalRowMutation, Boolean> {
3428

35-
private final SessionPoolMap<TableId, TableAsync> tables;
36-
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> authViews;
29+
private final SessionPoolMap sessionMap;
3730

38-
public CheckAndMutateRowShim(Client client) {
39-
// TODO: once the backend confirms the required permission for CheckAndMutateRow, if it is
40-
// PERMISSION_READ_WRITE, merge these pools with ReadModifyWriteRowShim's pools into a shared
41-
// ReadWriteSessionPools to halve session connections and unify circuit-breaker tracking.
42-
// CheckAndMutateRow reads (predicate filter) and conditionally writes, so it needs
43-
// read + write access on the session.
44-
tables =
45-
new SessionPoolMap<>(
46-
k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ_WRITE));
47-
authViews =
48-
new SessionPoolMap<>(
49-
k ->
50-
client.openAuthorizedViewAsync(
51-
k.getTableId(),
52-
k.getAuthorizedViewId(),
53-
OpenAuthorizedViewRequest.Permission.PERMISSION_READ_WRITE));
31+
public CheckAndMutateRowShim(SessionPoolMap sessionMap) {
32+
this.sessionMap = sessionMap;
5433
}
5534

5635
@Override
57-
public void close() throws IOException {
58-
tables.invalidateAll();
59-
authViews.invalidateAll();
60-
}
36+
public void close() {}
6137

6238
@Override
6339
public boolean supports(ConditionalRowMutation request) {
6440
TargetId targetId = request.getTargetId();
6541
SessionPool<?> pool;
6642
if (targetId instanceof TableId) {
67-
pool = tables.get((TableId) targetId).getSessionPool();
43+
pool = sessionMap.get((TableId) targetId, SessionKey.Permission.READ_WRITE).getSessionPool();
6844
} else if (targetId instanceof AuthorizedViewId) {
69-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
45+
pool =
46+
sessionMap
47+
.get((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE)
48+
.getSessionPool();
7049
} else {
7150
return false;
7251
}
@@ -80,13 +59,15 @@ public CompletableFuture<Boolean> call(ConditionalRowMutation request, Deadline
8059
SessionCheckAndMutateRowRequest innerReq = request.toSessionProto();
8160

8261
if (targetId instanceof TableId) {
83-
return tables.apply(
62+
return sessionMap.apply(
8463
(TableId) targetId,
64+
SessionKey.Permission.READ_WRITE,
8565
t -> t.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched()));
8666
}
8767
if (targetId instanceof AuthorizedViewId) {
88-
return authViews.apply(
68+
return sessionMap.apply(
8969
(AuthorizedViewId) targetId,
70+
SessionKey.Permission.READ_WRITE,
9071
v -> v.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched()));
9172
}
9273

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,15 +109,16 @@ public ApiFuture<RespT> futureCall(ReqT request, ApiCallContext context) {
109109
}
110110

111111
private boolean useExperimental(ReqT req) {
112-
if (!experimental.supports(req)) {
113-
return false;
114-
}
115-
116112
float ratio =
117113
configurationManager.getClientConfiguration().getSessionConfiguration().getSessionLoad();
118114
if (ratio == 0) {
119115
return false;
120116
}
117+
118+
if (!experimental.supports(req)) {
119+
return false;
120+
}
121+
121122
if (ratio != 1.0 && ThreadLocalRandom.current().nextFloat() > ratio) {
122123
return false;
123124
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java‎

Lines changed: 15 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -15,54 +15,35 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.internal.compat.ops;
1717

18-
import com.google.bigtable.v2.OpenAuthorizedViewRequest;
19-
import com.google.bigtable.v2.OpenTableRequest.Permission;
2018
import com.google.bigtable.v2.SessionMutateRowRequest;
21-
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
22-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
23-
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2419
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2520
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2621
import com.google.cloud.bigtable.data.v2.models.RowMutation;
2722
import com.google.cloud.bigtable.data.v2.models.TableId;
2823
import com.google.cloud.bigtable.data.v2.models.TargetId;
2924
import io.grpc.Deadline;
30-
import java.io.IOException;
3125
import java.util.concurrent.CompletableFuture;
3226

3327
public class MutateRowShim implements UnaryShim<RowMutation, Void> {
3428

35-
private final SessionPoolMap<TableId, TableAsync> tables;
36-
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> authViews;
29+
private final SessionPoolMap sessionMap;
3730

38-
public MutateRowShim(Client client) {
39-
tables =
40-
new SessionPoolMap<>(
41-
k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_WRITE));
42-
authViews =
43-
new SessionPoolMap<>(
44-
k ->
45-
client.openAuthorizedViewAsync(
46-
k.getTableId(),
47-
k.getAuthorizedViewId(),
48-
OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE));
31+
public MutateRowShim(SessionPoolMap sessionMap) {
32+
this.sessionMap = sessionMap;
4933
}
5034

5135
@Override
52-
public void close() throws IOException {
53-
tables.invalidateAll();
54-
authViews.invalidateAll();
55-
}
36+
public void close() {}
5637

5738
@Override
5839
public boolean supports(RowMutation request) {
5940
TargetId targetId = request.getTargetId();
6041
SessionPool<?> pool;
61-
// TODO: avoid double lookup
6242
if (targetId instanceof TableId) {
63-
pool = tables.get((TableId) targetId).getSessionPool();
43+
pool = sessionMap.get((TableId) targetId, SessionKey.Permission.WRITE).getSessionPool();
6444
} else if (targetId instanceof AuthorizedViewId) {
65-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
45+
pool =
46+
sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.WRITE).getSessionPool();
6647
} else {
6748
return false;
6849
}
@@ -76,12 +57,16 @@ public CompletableFuture<Void> call(RowMutation request, Deadline deadline) {
7657
SessionMutateRowRequest innerReq = request.toSessionProto();
7758

7859
if (targetId instanceof TableId) {
79-
return tables.apply(
80-
(TableId) targetId, t -> t.mutateRow(innerReq, deadline).thenApply(r -> null));
60+
return sessionMap.apply(
61+
(TableId) targetId,
62+
SessionKey.Permission.WRITE,
63+
t -> t.mutateRow(innerReq, deadline).thenApply(r -> null));
8164
}
8265
if (targetId instanceof AuthorizedViewId) {
83-
return authViews.apply(
84-
(AuthorizedViewId) targetId, v -> v.mutateRow(innerReq, deadline).thenApply(r -> null));
66+
return sessionMap.apply(
67+
(AuthorizedViewId) targetId,
68+
SessionKey.Permission.WRITE,
69+
v -> v.mutateRow(innerReq, deadline).thenApply(r -> null));
8570
}
8671

8772
CompletableFuture<Void> f = new CompletableFuture<>();

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java‎

Lines changed: 17 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,8 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.internal.compat.ops;
1717

18-
import com.google.bigtable.v2.OpenAuthorizedViewRequest;
19-
import com.google.bigtable.v2.OpenTableRequest.Permission;
2018
import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
2119
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
22-
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
23-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
24-
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2520
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2621
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2722
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -33,38 +28,26 @@
3328
public class ReadModifyWriteRowShim
3429
implements UnaryShim<ReadModifyWriteRow, SessionReadModifyWriteRowResponse> {
3530

36-
private final SessionPoolMap<TableId, TableAsync> tables;
37-
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> authViews;
31+
private final SessionPoolMap sessionMap;
3832

39-
public ReadModifyWriteRowShim(Client client) {
40-
// ReadModifyWriteRow reads and modifies cells and returns the row, so it needs
41-
// read + write access on the session.
42-
tables =
43-
new SessionPoolMap<>(
44-
k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ_WRITE));
45-
authViews =
46-
new SessionPoolMap<>(
47-
k ->
48-
client.openAuthorizedViewAsync(
49-
k.getTableId(),
50-
k.getAuthorizedViewId(),
51-
OpenAuthorizedViewRequest.Permission.PERMISSION_READ_WRITE));
33+
public ReadModifyWriteRowShim(SessionPoolMap sessionMap) {
34+
this.sessionMap = sessionMap;
5235
}
5336

5437
@Override
55-
public void close() {
56-
tables.invalidateAll();
57-
authViews.invalidateAll();
58-
}
38+
public void close() {}
5939

6040
@Override
6141
public boolean supports(ReadModifyWriteRow request) {
6242
TargetId targetId = request.getTargetId();
6343
SessionPool<?> pool;
6444
if (targetId instanceof TableId) {
65-
pool = tables.get((TableId) targetId).getSessionPool();
45+
pool = sessionMap.get((TableId) targetId, SessionKey.Permission.READ_WRITE).getSessionPool();
6646
} else if (targetId instanceof AuthorizedViewId) {
67-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
47+
pool =
48+
sessionMap
49+
.get((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE)
50+
.getSessionPool();
6851
} else {
6952
return false;
7053
}
@@ -78,11 +61,16 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> call(
7861
SessionReadModifyWriteRowRequest innerReq = request.toSessionProto();
7962

8063
if (targetId instanceof TableId) {
81-
return tables.apply((TableId) targetId, t -> t.readModifyWriteRow(innerReq, deadline));
64+
return sessionMap.apply(
65+
(TableId) targetId,
66+
SessionKey.Permission.READ_WRITE,
67+
t -> t.readModifyWriteRow(innerReq, deadline));
8268
}
8369
if (targetId instanceof AuthorizedViewId) {
84-
return authViews.apply(
85-
(AuthorizedViewId) targetId, v -> v.readModifyWriteRow(innerReq, deadline));
70+
return sessionMap.apply(
71+
(AuthorizedViewId) targetId,
72+
SessionKey.Permission.READ_WRITE,
73+
v -> v.readModifyWriteRow(innerReq, deadline));
8674
}
8775

8876
CompletableFuture<SessionReadModifyWriteRowResponse> f = new CompletableFuture<>();

0 commit comments

Comments
 (0)