Skip to content

refactor(bigtable): make sessionMap shared by permissions #14662

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 1 commit into
base: main
Choose a base branch
from
+487 −212
Open

refactor(bigtable): make sessionMap shared by permissions #14662

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions ...ud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -85,6 +86,7 @@ public class ShimImpl implements Shim {
private final Resource<ClientConfigurationManager> configManagerResource;
private final Client client;
private final DebugTagTracer debugTagTracer;
private final SessionPoolMap sessionMap;

private final ReadRowShim readRowShim;
private final MutateRowShim mutateRowShim;
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -358,6 +361,7 @@ private static boolean checkDirectAccessAvailable(

@Override
public void close() {
sessionMap.invalidateAll();
client.close();
configManagerResource.close();
}
Expand Down
3 changes: 1 addition & 2 deletions ...-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Util.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -46,7 +45,7 @@ public static <K, V extends Closeable> LoadingCache<K, V> createSessionMap(
V value = notification.getValue();
//noinspection DataFlowIssue
value.close();
} catch (IOException e) {
} catch (Exception e) {
// TODO log
}
})
Expand Down
45 changes: 13 additions & 32 deletions ...ain/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/CheckAndMutateRowShim.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ConditionalRowMutation, Boolean> {

private final SessionPoolMap<TableId, TableAsync> tables;
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> 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;
Comment thread
mutianf marked this conversation as resolved.
}
Expand All @@ -80,13 +59,15 @@ public CompletableFuture<Boolean> 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()));
}

Expand Down
9 changes: 5 additions & 4 deletions ...in/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,15 +109,16 @@ public ApiFuture<RespT> 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;
}
Expand Down
45 changes: 15 additions & 30 deletions ...le/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<RowMutation, Void> {

private final SessionPoolMap<TableId, TableAsync> tables;
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> 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 {
Comment on lines 41 to 47

@gemini-code-assist gemini-code-assist Bot Oct 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The supports() method should use sessionMap.getIfPresent() instead of sessionMap.get() to avoid eagerly loading/opening the session pool during support checks. Eagerly loading the pool defeats the lazy-loading design and can cause unnecessary connection overhead or unexpected exceptions during inspection.

    SessionPool<?> pool;
    if (targetId instanceof TableId) {
      com.google.cloud.bigtable.data.v2.internal.api.TableAsync table =
          sessionMap.getIfPresent((TableId) targetId, SessionKey.Permission.WRITE);
      if (table == null) {
        return true;
      }
      pool = table.getSessionPool();
    } else if (targetId instanceof AuthorizedViewId) {
      com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync view =
          sessionMap.getIfPresent((AuthorizedViewId) targetId, SessionKey.Permission.WRITE);
      if (view == null) {
        return true;
      }
      pool = view.getSessionPool();
    } else {

return false;
}
Expand All @@ -76,12 +57,16 @@ public CompletableFuture<Void> 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<Void> f = new CompletableFuture<>();
Expand Down
46 changes: 17 additions & 29 deletions ...in/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadModifyWriteRowShim.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -33,38 +28,26 @@
public class ReadModifyWriteRowShim
implements UnaryShim<ReadModifyWriteRow, SessionReadModifyWriteRowResponse> {

private final SessionPoolMap<TableId, TableAsync> tables;
private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> 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 {
Comment on lines 43 to 51

@gemini-code-assist gemini-code-assist Bot Oct 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The supports() method should use sessionMap.getIfPresent() instead of sessionMap.get() to avoid eagerly loading/opening the session pool during support checks. Eagerly loading the pool defeats the lazy-loading design and can cause unnecessary connection overhead or unexpected exceptions during inspection.

    SessionPool<?> pool;
    if (targetId instanceof TableId) {
      com.google.cloud.bigtable.data.v2.internal.api.TableAsync table =
          sessionMap.getIfPresent((TableId) targetId, SessionKey.Permission.READ_WRITE);
      if (table == null) {
        return true;
      }
      pool = table.getSessionPool();
    } else if (targetId instanceof AuthorizedViewId) {
      com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync view =
          sessionMap.getIfPresent((AuthorizedViewId) targetId, SessionKey.Permission.READ_WRITE);
      if (view == null) {
        return true;
      }
      pool = view.getSessionPool();
    } else {

return false;
}
Expand All @@ -78,11 +61,16 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> 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<SessionReadModifyWriteRowResponse> f = new CompletableFuture<>();
Expand Down
Loading
Loading