Skip to content

Commit fc9b8d8

Browse files
mutianfmutianf
committed
refactor(bigtable): introduce SessionKey/SessionPermission; merge SessionMap into SessionPoolMap
Add SessionPermission (READ/WRITE/READ_WRITE) and SessionKey<IdT> to pair a target ID with a permission level, then absorb the three-map registry logic directly into SessionPoolMap instead of a separate class. SessionPoolMap now holds one LoadingCache per ID type (TableId / AuthorizedViewId / MaterializedViewId), keyed by SessionKey<IdT>. Two shim operations requesting the same (id, permission) — e.g. a future CheckAndMutateRowShim and ReadModifyWriteRowShim both using READ_WRITE on the same table — will resolve to the same handle and therefore the same session pool. Fix supports() to use getIfPresent() instead of get(). Previously supports() called the loading get(), which would: 1. Open a session pool as a side effect of a cheap read-only check. 2. Throw IllegalStateException (Client is closed) unchecked through DivertingUnaryCallable.useExperimental(), bypassing the normal failed-future error surface. With getIfPresent(), a null return (no pool yet) is treated as "optimistically supported": the pool opens lazily in call(), and if that fails (closed client, etc.) apply() surfaces it as a failed future. ShimImpl creates a single shared SessionPoolMap and passes it to all shims; each shim declares its required SessionPermission internally.
1 parent 7dd2b37 commit fc9b8d8

10 files changed

Lines changed: 594 additions & 203 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: 16 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -15,58 +15,42 @@
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;
2119
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
22-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
2320
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2421
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2522
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2623
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
2724
import com.google.cloud.bigtable.data.v2.models.TableId;
2825
import com.google.cloud.bigtable.data.v2.models.TargetId;
2926
import io.grpc.Deadline;
30-
import java.io.IOException;
3127
import java.util.concurrent.CompletableFuture;
3228

3329
public class CheckAndMutateRowShim implements UnaryShim<ConditionalRowMutation, Boolean> {
3430

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

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));
33+
public CheckAndMutateRowShim(SessionPoolMap sessionMap) {
34+
this.sessionMap = sessionMap;
5435
}
5536

5637
@Override
57-
public void close() throws IOException {
58-
tables.invalidateAll();
59-
authViews.invalidateAll();
60-
}
38+
public void close() {}
6139

6240
@Override
6341
public boolean supports(ConditionalRowMutation request) {
6442
TargetId targetId = request.getTargetId();
6543
SessionPool<?> pool;
6644
if (targetId instanceof TableId) {
67-
pool = tables.get((TableId) targetId).getSessionPool();
45+
TableAsync handle =
46+
sessionMap.getIfPresent((TableId) targetId, SessionKey.Permission.READ_WRITE);
47+
if (handle == null) return true;
48+
pool = handle.getSessionPool();
6849
} else if (targetId instanceof AuthorizedViewId) {
69-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
50+
AuthorizedViewAsync handle =
51+
sessionMap.getIfPresent((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE);
52+
if (handle == null) return true;
53+
pool = handle.getSessionPool();
7054
} else {
7155
return false;
7256
}
@@ -80,13 +64,15 @@ public CompletableFuture<Boolean> call(ConditionalRowMutation request, Deadline
8064
SessionCheckAndMutateRowRequest innerReq = request.toSessionProto();
8165

8266
if (targetId instanceof TableId) {
83-
return tables.apply(
67+
return sessionMap.apply(
8468
(TableId) targetId,
69+
SessionKey.Permission.READ_WRITE,
8570
t -> t.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched()));
8671
}
8772
if (targetId instanceof AuthorizedViewId) {
88-
return authViews.apply(
73+
return sessionMap.apply(
8974
(AuthorizedViewId) targetId,
75+
SessionKey.Permission.READ_WRITE,
9076
v -> v.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched()));
9177
}
9278

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

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -15,54 +15,41 @@
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;
2119
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
22-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
2320
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2421
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2522
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2623
import com.google.cloud.bigtable.data.v2.models.RowMutation;
2724
import com.google.cloud.bigtable.data.v2.models.TableId;
2825
import com.google.cloud.bigtable.data.v2.models.TargetId;
2926
import io.grpc.Deadline;
30-
import java.io.IOException;
3127
import java.util.concurrent.CompletableFuture;
3228

3329
public class MutateRowShim implements UnaryShim<RowMutation, Void> {
3430

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

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));
33+
public MutateRowShim(SessionPoolMap sessionMap) {
34+
this.sessionMap = sessionMap;
4935
}
5036

5137
@Override
52-
public void close() throws IOException {
53-
tables.invalidateAll();
54-
authViews.invalidateAll();
55-
}
38+
public void close() {}
5639

5740
@Override
5841
public boolean supports(RowMutation request) {
5942
TargetId targetId = request.getTargetId();
6043
SessionPool<?> pool;
61-
// TODO: avoid double lookup
6244
if (targetId instanceof TableId) {
63-
pool = tables.get((TableId) targetId).getSessionPool();
45+
TableAsync handle = sessionMap.getIfPresent((TableId) targetId, SessionKey.Permission.WRITE);
46+
if (handle == null) return true;
47+
pool = handle.getSessionPool();
6448
} else if (targetId instanceof AuthorizedViewId) {
65-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
49+
AuthorizedViewAsync handle =
50+
sessionMap.getIfPresent((AuthorizedViewId) targetId, SessionKey.Permission.WRITE);
51+
if (handle == null) return true;
52+
pool = handle.getSessionPool();
6653
} else {
6754
return false;
6855
}
@@ -76,12 +63,16 @@ public CompletableFuture<Void> call(RowMutation request, Deadline deadline) {
7663
SessionMutateRowRequest innerReq = request.toSessionProto();
7764

7865
if (targetId instanceof TableId) {
79-
return tables.apply(
80-
(TableId) targetId, t -> t.mutateRow(innerReq, deadline).thenApply(r -> null));
66+
return sessionMap.apply(
67+
(TableId) targetId,
68+
SessionKey.Permission.WRITE,
69+
t -> t.mutateRow(innerReq, deadline).thenApply(r -> null));
8170
}
8271
if (targetId instanceof AuthorizedViewId) {
83-
return authViews.apply(
84-
(AuthorizedViewId) targetId, v -> v.mutateRow(innerReq, deadline).thenApply(r -> null));
72+
return sessionMap.apply(
73+
(AuthorizedViewId) targetId,
74+
SessionKey.Permission.WRITE,
75+
v -> v.mutateRow(innerReq, deadline).thenApply(r -> null));
8576
}
8677

8778
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: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,9 @@
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;
2220
import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync;
23-
import com.google.cloud.bigtable.data.v2.internal.api.Client;
2421
import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
2522
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
2623
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
@@ -33,38 +30,29 @@
3330
public class ReadModifyWriteRowShim
3431
implements UnaryShim<ReadModifyWriteRow, SessionReadModifyWriteRowResponse> {
3532

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

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));
35+
public ReadModifyWriteRowShim(SessionPoolMap sessionMap) {
36+
this.sessionMap = sessionMap;
5237
}
5338

5439
@Override
55-
public void close() {
56-
tables.invalidateAll();
57-
authViews.invalidateAll();
58-
}
40+
public void close() {}
5941

6042
@Override
6143
public boolean supports(ReadModifyWriteRow request) {
6244
TargetId targetId = request.getTargetId();
6345
SessionPool<?> pool;
6446
if (targetId instanceof TableId) {
65-
pool = tables.get((TableId) targetId).getSessionPool();
47+
TableAsync handle =
48+
sessionMap.getIfPresent((TableId) targetId, SessionKey.Permission.READ_WRITE);
49+
if (handle == null) return true;
50+
pool = handle.getSessionPool();
6651
} else if (targetId instanceof AuthorizedViewId) {
67-
pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
52+
AuthorizedViewAsync handle =
53+
sessionMap.getIfPresent((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE);
54+
if (handle == null) return true;
55+
pool = handle.getSessionPool();
6856
} else {
6957
return false;
7058
}
@@ -78,11 +66,16 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> call(
7866
SessionReadModifyWriteRowRequest innerReq = request.toSessionProto();
7967

8068
if (targetId instanceof TableId) {
81-
return tables.apply((TableId) targetId, t -> t.readModifyWriteRow(innerReq, deadline));
69+
return sessionMap.apply(
70+
(TableId) targetId,
71+
SessionKey.Permission.READ_WRITE,
72+
t -> t.readModifyWriteRow(innerReq, deadline));
8273
}
8374
if (targetId instanceof AuthorizedViewId) {
84-
return authViews.apply(
85-
(AuthorizedViewId) targetId, v -> v.readModifyWriteRow(innerReq, deadline));
75+
return sessionMap.apply(
76+
(AuthorizedViewId) targetId,
77+
SessionKey.Permission.READ_WRITE,
78+
v -> v.readModifyWriteRow(innerReq, deadline));
8679
}
8780

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

0 commit comments

Comments
 (0)