From 7d54285fc01ab344f43d11ebe0bf5b2a6536ec4d Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Fri, 26 Jun 2026 23:40:00 +0000 Subject: [PATCH 1/4] refactor: replace NettyWheelTimer with JDK ScheduledExecutorTimer Drops the shaded-netty internal reach-in (HashedWheelTimer) in favor of ScheduledExecutorService. BigtableTimer.newTimeout now takes a caller-supplied Executor so timer callbacks hop off the scheduler thread to the appropriate pool per call site. Benchmark (jetstream100, n=3 per scenario, c32/c100/c500): throughput within 1-2% of the prior NettyWheelTimer impl; CPU/op within 1-3%; p50, p99, p99.9 all within noise. --- .../v2/internal/api/AuthorizedViewAsync.java | 2 + .../bigtable/data/v2/internal/api/Client.java | 9 +- .../internal/api/MaterializedViewAsync.java | 2 + .../data/v2/internal/api/TableAsync.java | 2 + .../data/v2/internal/api/TableBase.java | 4 +- .../v2/internal/middleware/RetryingVRpc.java | 15 +-- .../v2/internal/session/BigtableTimer.java | 27 ++-- .../v2/internal/session/NettyWheelTimer.java | 116 ----------------- .../session/ScheduledExecutorTimer.java | 114 ++++++++++++++++ .../data/v2/internal/session/SessionImpl.java | 3 +- .../v2/internal/session/SessionPoolImpl.java | 43 +++++- .../internal/csm/tracers/VRpcTracerTest.java | 6 +- .../internal/middleware/RetryingVRpcTest.java | 6 +- .../session/ScheduledExecutorTimerTest.java | 122 ++++++++++++++++++ .../v2/internal/session/SessionImplTest.java | 9 +- .../internal/session/SessionPoolImplTest.java | 31 +++-- .../v2/internal/session/WatchdogTest.java | 5 +- 17 files changed, 338 insertions(+), 178 deletions(-) delete mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java index adb6b2e2e0c0..490bbf06b8a1 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java @@ -50,6 +50,7 @@ static AuthorizedViewAsync createAndStart( Permission permission, Metrics metrics, BigtableTimer timer, + Executor backgroundExecutor, Executor userCallbackExecutor) { AuthorizedViewName viewName = @@ -81,6 +82,7 @@ static AuthorizedViewAsync createAndStart( viewName.toString(), metrics, timer, + backgroundExecutor, userCallbackExecutor); return new AuthorizedViewAsync(base); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index ec4cbd9a3f06..4ca588403702 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -34,7 +34,7 @@ import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; -import com.google.cloud.bigtable.data.v2.internal.session.NettyWheelTimer; +import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.util.concurrent.ThreadFactoryBuilder; @@ -205,9 +205,7 @@ public Client( this.configManager = configManager; this.backgroundExecutor = bgExecutor; this.userCallbackExecutor = userCallbackExecutor; - // Timer's tick thread dispatches bodies onto backgroundExecutor — tick-thread-blocking work - // (anything that takes a pool lock) gets handed off there instead of stalling the wheel. - this.sessionTimer = new NettyWheelTimer("bigtable-session-timer", bgExecutor.get()); + this.sessionTimer = new ScheduledExecutorTimer("bigtable-session-timer"); defaultCallOptions = CallOptions.DEFAULT; @@ -311,6 +309,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) { permission, metrics.get(), sessionTimer, + backgroundExecutor.get(), userCallbackExecutor.get()); sessionPools.add(tableAsync.getSessionPool()); return tableAsync; @@ -335,6 +334,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync( permission, metrics.get(), sessionTimer, + backgroundExecutor.get(), userCallbackExecutor.get()); sessionPools.add(viewAsync.getSessionPool()); return viewAsync; @@ -358,6 +358,7 @@ public MaterializedViewAsync openMaterializedViewAsync( permission, metrics.get(), sessionTimer, + backgroundExecutor.get(), userCallbackExecutor.get()); sessionPools.add(viewAsync.getSessionPool()); return viewAsync; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java index eb25b1d858dd..98d9d31990f8 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java @@ -46,6 +46,7 @@ public static MaterializedViewAsync createAndStart( OpenMaterializedViewRequest.Permission permission, Metrics metrics, BigtableTimer timer, + Executor backgroundExecutor, Executor userCallbackExecutor) { MaterializedViewName viewName = @@ -76,6 +77,7 @@ public static MaterializedViewAsync createAndStart( viewId, metrics, timer, + backgroundExecutor, userCallbackExecutor); return new MaterializedViewAsync(base); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java index 8bb16c51c201..049d45df8498 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java @@ -49,6 +49,7 @@ public static TableAsync createAndStart( Permission permission, Metrics metrics, BigtableTimer timer, + Executor backgroundExecutor, Executor userCallbackExecutor) { TableName tableName = @@ -79,6 +80,7 @@ public static TableAsync createAndStart( tableId, metrics, timer, + backgroundExecutor, userCallbackExecutor); return new TableAsync(base); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java index a3b4a535b0d9..5565e9687a3d 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java @@ -63,6 +63,7 @@ static TableBase createAndStart( String sessionPoolName, Metrics metrics, BigtableTimer timer, + Executor backgroundExecutor, Executor userCallbackExecutor) { SessionPool sessionPool = @@ -75,7 +76,8 @@ static TableBase createAndStart( callOptions, sessionDescriptor, sessionPoolName, - timer); + timer, + backgroundExecutor); sessionPool.start(openReq, new Metadata()); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java index d11c509798d4..434577bc138b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java @@ -329,19 +329,14 @@ public void onStart() { try { stopHook = timer.onStop(this::onTimerStopping); // Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at - // the moment the body runs, not just while the dispatcher is invoking the outer task. + // the moment the body runs on the op executor. future = timer.newTimeout( () -> - context - .getExecutor() - .execute( - () -> - grpcContext - .wrap( - () -> - otelContext.wrap(() -> onStateChange(new Idle())).run()) - .run()), + grpcContext + .wrap(() -> otelContext.wrap(() -> onStateChange(new Idle())).run()) + .run(), + context.getExecutor(), Durations.toMillis(retryDelay), TimeUnit.MILLISECONDS); } catch (IllegalStateException e) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java index e3edd34fd606..788ff433c3c2 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java @@ -15,6 +15,7 @@ */ package com.google.cloud.bigtable.data.v2.internal.session; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; /** @@ -22,28 +23,24 @@ * approximate, low-resolution times. Backed by a hashed wheel: O(1) insert and O(1) cancel, * regardless of how many pending timeouts the wheel holds. * - *

{@link #newTimeout} runs the callback on the timer's bundled dispatch executor — callers do - * not have to wrap their bodies in {@code executor.execute(...)} to stay off the tick thread. This - * is the default and is correct for any callback that takes a lock or does real work. - * - *

TODO: once later refactor steps introduce per-op / per-session dispatchers (e.g. {@code - * SerializingExecutor} or {@code SynchronizationContext}), add a {@code newTimeoutOnTickThread} - * variant so callers with their own dispatcher can skip the wasted hop through the bundled - * executor. - * - *

This is a thin abstraction over a single concrete implementation today (see {@code - * NettyWheelTimer}). It exists so the implementation can be swapped after benchmarking establishes - * a baseline. + *

Each {@link #newTimeout} call carries the {@link Executor} that should run the task body. + * Tasks never run on the timer's tick thread (unless the caller passes {@code directExecutor()}; + * see warning below). Passing the executor per-call avoids the dispatcher hop that a bundled + * default would impose on callers who already trampoline onto their own executor. */ public interface BigtableTimer { /** - * Schedules {@code task} to run after {@code delay}. The task body runs on the timer's bundled - * dispatch executor, not on the tick thread, so it is safe to take locks or do bounded work. + * Schedules {@code task} to run after {@code delay}. The task body is handed to {@code + * executor.execute(task)} — pass the executor where this task should ultimately run. * *

The returned handle can be used to cancel the task; cancel is O(1) and does not leave the * entry in any heap. + * + *

Warning: passing {@code MoreExecutors.directExecutor()} runs {@code task} inline on + * the timer's tick thread. The task must be trivial and non-blocking — anything more will stall + * every other scheduled timeout on the wheel. */ - Timeout newTimeout(Runnable task, long delay, TimeUnit unit); + Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit); /** * Releases the tick thread and discards any pending timeouts. Idempotent. After {@code stop()}, diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java deleted file mode 100644 index 16ffd0127b10..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java +++ /dev/null @@ -1,116 +0,0 @@ -/* - * 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.session; - -import com.google.common.util.concurrent.ThreadFactoryBuilder; -import io.grpc.netty.shaded.io.netty.util.HashedWheelTimer; -import java.util.HashSet; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import java.util.logging.Level; -import java.util.logging.Logger; - -/** - * {@link BigtableTimer} backed by Netty's {@code HashedWheelTimer}, accessed via the shaded copy - * inside {@code grpc-netty-shaded}. We depend on the shaded class so the library does not pull in - * an additional {@code io.netty:netty-common} artifact. - * - *

Temporary: once threading-refactor benchmarks establish a baseline, this should be replaced - * with an in-tree implementation that does not reach into gRPC's shaded internals. - */ -public final class NettyWheelTimer implements BigtableTimer { - private static final Logger LOG = Logger.getLogger(NettyWheelTimer.class.getName()); - - // 10 ms tick × 512 buckets ≈ 5 s per rotation. Heartbeat (100 ms), deadlines (sub-second to - // seconds), and watchdog (5 min) all sit comfortably inside this resolution. - private static final long TICK_DURATION_MS = 10; - private static final int TICKS_PER_WHEEL = 512; - - private final HashedWheelTimer delegate; - private final Executor dispatcher; - - // ConcurrentHashMap-backed Set so onStop/Registration.unregister can run from any thread without - // blocking newTimeout. Stop drains it once, then refuses further registrations. - private final Set stopHooks = ConcurrentHashMap.newKeySet(); - private volatile boolean stopped = false; - - public NettyWheelTimer(String name, Executor dispatcher) { - this.dispatcher = dispatcher; - this.delegate = - new HashedWheelTimer( - new ThreadFactoryBuilder().setNameFormat(name + "-%d").setDaemon(true).build(), - TICK_DURATION_MS, - TimeUnit.MILLISECONDS, - TICKS_PER_WHEEL); - } - - @Override - public Timeout newTimeout(Runnable task, long delay, TimeUnit unit) { - if (stopped) { - throw new IllegalStateException("timer stopped"); - } - return new TimeoutHandle(delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit)); - } - - @Override - public Registration onStop(Runnable hook) { - if (stopped) { - throw new IllegalStateException("timer stopped"); - } - stopHooks.add(hook); - return () -> stopHooks.remove(hook); - } - - @Override - public void stop() { - if (stopped) { - return; - } - stopped = true; - // Snapshot then clear so hooks can no longer be unregistered while we iterate (and so a hook - // that re-enters onStop sees stopped=true and fails fast). - Set hooks = new HashSet<>(stopHooks); - stopHooks.clear(); - for (Runnable hook : hooks) { - try { - hook.run(); - } catch (Throwable t) { - LOG.log(Level.WARNING, "stop hook threw; continuing", t); - } - } - delegate.stop(); - } - - private static final class TimeoutHandle implements Timeout { - private final io.grpc.netty.shaded.io.netty.util.Timeout nettyTimeout; - - TimeoutHandle(io.grpc.netty.shaded.io.netty.util.Timeout nettyTimeout) { - this.nettyTimeout = nettyTimeout; - } - - @Override - public boolean cancel() { - return nettyTimeout.cancel(); - } - - @Override - public boolean isCancelled() { - return nettyTimeout.isCancelled(); - } - } -} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java new file mode 100644 index 000000000000..25ce016b73d3 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java @@ -0,0 +1,114 @@ +/* + * 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.session; + +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import java.util.HashSet; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * {@link BigtableTimer} backed by a single-threaded {@link ScheduledExecutorService}. + * + *

Each {@code newTimeout} schedules a wrapper that hops to the caller-supplied {@link Executor}. + * Cancel delegates to {@link ScheduledFuture#cancel(boolean)} with {@code mayInterruptIfRunning = + * false}. + */ +public final class ScheduledExecutorTimer implements BigtableTimer { + private static final Logger LOG = Logger.getLogger(ScheduledExecutorTimer.class.getName()); + + private final ScheduledExecutorService scheduler; + private final Set stopHooks = ConcurrentHashMap.newKeySet(); + + private volatile boolean stopped = false; + + public ScheduledExecutorTimer(String name) { + this.scheduler = + Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryBuilder().setNameFormat(name + "-%d").setDaemon(true).build()); + } + + @Override + public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + ScheduledFuture future = + scheduler.schedule( + () -> { + try { + executor.execute(task); + } catch (Throwable t) { + LOG.log(Level.WARNING, "executor.execute threw; dropping task", t); + } + }, + delay, + unit); + return new FutureTimeout(future); + } + + @Override + public Registration onStop(Runnable hook) { + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + stopHooks.add(hook); + return () -> stopHooks.remove(hook); + } + + @Override + public void stop() { + if (stopped) { + return; + } + stopped = true; + Set hooks = new HashSet<>(stopHooks); + stopHooks.clear(); + for (Runnable hook : hooks) { + try { + hook.run(); + } catch (Throwable t) { + LOG.log(Level.WARNING, "stop hook threw; continuing", t); + } + } + scheduler.shutdownNow(); + } + + private static final class FutureTimeout implements Timeout { + private final ScheduledFuture future; + + FutureTimeout(ScheduledFuture future) { + this.future = future; + } + + @Override + public boolean cancel() { + return future.cancel(false); + } + + @Override + public boolean isCancelled() { + return future.isCancelled(); + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java index 1915edf5306e..50ef559fc970 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java @@ -455,7 +455,8 @@ public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable private void scheduleHeartbeatCheck() { heartbeatTimeout = timer.newTimeout( - () -> sessionSyncContext.execute(this::checkHeartbeat), + this::checkHeartbeat, + sessionSyncContext, HEARTBEAT_CHECK_INTERVAL.toMillis(), TimeUnit.MILLISECONDS); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java index d3e65975a394..50542c7f5fdb 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java @@ -66,6 +66,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.logging.Level; @@ -118,6 +119,11 @@ private enum PoolState { // this pool. One tick thread per Client (owned by Client); O(1) insert / O(1) cancel. private final BigtableTimer timer; + // Executor for timer-scheduled pool maintenance bodies (AFE prune, retry-create-session, + // deadline monitor, watchdog tick). The timer's tick thread hands work here so it doesn't + // run pool-lock-holding work inline. + private final Executor backgroundExecutor; + @GuardedBy("this") private int consecutiveFailures = 0; @@ -167,7 +173,8 @@ public SessionPoolImpl( CallOptions callOptions, SessionDescriptor sessionDescriptor, String name, - BigtableTimer timer) { + BigtableTimer timer, + Executor backgroundExecutor) { this( metrics, featureFlags, @@ -178,6 +185,7 @@ public SessionPoolImpl( sessionDescriptor, name, timer, + backgroundExecutor, createInitialBudget(configManager.getClientConfiguration())); } @@ -194,6 +202,7 @@ public SessionPoolImpl( SessionDescriptor sessionDescriptor, String name, BigtableTimer timer, + Executor backgroundExecutor, SessionCreationBudget budget) { this.metrics = metrics; this.featureFlags = featureFlags; @@ -204,6 +213,7 @@ public SessionPoolImpl( // Timer is owned by the caller (typically Client) and shared across pools — do NOT stop it // in close(). this.timer = timer; + this.backgroundExecutor = backgroundExecutor; sessions = new SessionList(); LoadBalancingOptions lbOptions = @@ -225,7 +235,9 @@ public SessionPoolImpl( debugTagTracer = metrics.getDebugTagTracer(); // Watchdog checks for sessions in WAIT_SERVER_CLOSE state and runs every 5 minutes - watchdog = new Watchdog(this, timer, Duration.ofMinutes(5), sessions, debugTagTracer); + watchdog = + new Watchdog( + this, timer, backgroundExecutor, Duration.ofMinutes(5), sessions, debugTagTracer); // Heartbeat monitoring is now done per-session via SessionImpl.scheduleHeartbeatCheck. scheduleNextAfePrune(); @@ -252,6 +264,7 @@ private void scheduleNextAfePrune() { afeListPruneTimeout = timer.newTimeout( this::runAfePruneAndReschedule, + backgroundExecutor, SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(), TimeUnit.MILLISECONDS); } @@ -461,6 +474,7 @@ private void maybeScheduleCreateSessionRetry() { } } }, + backgroundExecutor, retryIntervalMs, TimeUnit.MILLISECONDS); } @@ -814,14 +828,14 @@ private VRpcListener getListener() { } private BigtableTimer.Timeout monitorDeadline(Deadline deadline) { - // Body runs on the timer's bundled dispatcher (off the tick thread). - // onlyCancelPendingCall=true avoids racing with a user cancel that already attached a real - // call. + // Body runs on backgroundExecutor (off the tick thread). onlyCancelPendingCall=true avoids + // racing with a user cancel that already attached a real call. return timer.newTimeout( () -> cancel( Status.DEADLINE_EXCEEDED.withDescription("Deadline exceeded waiting for session"), true), + backgroundExecutor, deadline.timeRemaining(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS); } @@ -832,6 +846,7 @@ static class Watchdog implements Runnable { private final Object lock; private final BigtableTimer timer; + private final Executor backgroundExecutor; private final Duration interval; private final SessionList sessions; private final Clock clock; @@ -854,22 +869,32 @@ static class Watchdog implements Runnable { public Watchdog( Object lock, BigtableTimer timer, + Executor backgroundExecutor, Duration interval, SessionList sessionList, DebugTagTracer debugTagTracer) { - this(lock, timer, interval, sessionList, debugTagTracer, Clock.systemUTC()); + this( + lock, + timer, + backgroundExecutor, + interval, + sessionList, + debugTagTracer, + Clock.systemUTC()); } @VisibleForTesting Watchdog( Object lock, BigtableTimer timer, + Executor backgroundExecutor, Duration interval, SessionList sessionList, DebugTagTracer debugTagTracer, Clock clock) { this.lock = lock; this.timer = timer; + this.backgroundExecutor = backgroundExecutor; this.interval = interval; this.sessions = sessionList; this.debugTagTracer = debugTagTracer; @@ -885,7 +910,11 @@ private void scheduleNext() { synchronized (scheduleLock) { if (watchdogClosed) return; currentTimeout = - timer.newTimeout(this::runAndReschedule, interval.toMillis(), TimeUnit.MILLISECONDS); + timer.newTimeout( + this::runAndReschedule, + backgroundExecutor, + interval.toMillis(), + TimeUnit.MILLISECONDS); } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java index 4602c0941822..1ecc6e88b0d3 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java @@ -46,10 +46,10 @@ import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; import com.google.cloud.bigtable.data.v2.internal.session.FakeDescriptor; -import com.google.cloud.bigtable.data.v2.internal.session.NettyWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.Session; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; +import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; @@ -118,9 +118,7 @@ public class VRpcTracerTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = - new NettyWheelTimer( - "vrpc-tracer-test", com.google.common.util.concurrent.MoreExecutors.directExecutor()); + timer = new ScheduledExecutorTimer("vrpc-tracer-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java index 6776aef7ef48..e89f9de108e5 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java @@ -39,9 +39,9 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; import com.google.cloud.bigtable.data.v2.internal.session.FakeDescriptor; -import com.google.cloud.bigtable.data.v2.internal.session.NettyWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; +import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; @@ -91,9 +91,7 @@ public class RetryingVRpcTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = - new NettyWheelTimer( - "retrying-vrpc-test", com.google.common.util.concurrent.MoreExecutors.directExecutor()); + timer = new ScheduledExecutorTimer("retrying-vrpc-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java new file mode 100644 index 000000000000..189a00903748 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java @@ -0,0 +1,122 @@ +/* + * 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.session; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer.Registration; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +@Timeout(value = 10, unit = TimeUnit.SECONDS) +class ScheduledExecutorTimerTest { + private static final Executor DIRECT = MoreExecutors.directExecutor(); + + private ScheduledExecutorTimer timer; + + @AfterEach + void tearDown() { + if (timer != null) { + timer.stop(); + } + } + + @Test + void schedule_fires_after_delay() throws Exception { + timer = new ScheduledExecutorTimer("ses-test"); + CountDownLatch latch = new CountDownLatch(1); + + timer.newTimeout(latch::countDown, DIRECT, 10, TimeUnit.MILLISECONDS); + + assertThat(latch.await(1, TimeUnit.SECONDS)).isTrue(); + } + + @Test + void cancel_before_fire_prevents_execution() throws Exception { + timer = new ScheduledExecutorTimer("ses-test"); + AtomicInteger fires = new AtomicInteger(); + + BigtableTimer.Timeout t = + timer.newTimeout(fires::incrementAndGet, DIRECT, 100, TimeUnit.MILLISECONDS); + assertThat(t.cancel()).isTrue(); + assertThat(t.isCancelled()).isTrue(); + + Thread.sleep(200); + assertThat(fires.get()).isEqualTo(0); + } + + @Test + void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() { + timer = new ScheduledExecutorTimer("ses-test"); + Thread caller = Thread.currentThread(); + List ran = new CopyOnWriteArrayList<>(); + + timer.onStop(() -> ran.add(Thread.currentThread())); + timer.onStop(() -> ran.add(Thread.currentThread())); + + timer.stop(); + + assertThat(ran).hasSize(2); + for (Thread t : ran) { + assertThat(t).isSameInstanceAs(caller); + } + assertThrows( + IllegalStateException.class, + () -> timer.newTimeout(() -> {}, DIRECT, 1, TimeUnit.MILLISECONDS)); + timer = null; + } + + @Test + void unregister_removes_hook() { + timer = new ScheduledExecutorTimer("ses-test"); + AtomicInteger fires = new AtomicInteger(); + + Registration reg = timer.onStop(fires::incrementAndGet); + reg.unregister(); + + timer.stop(); + timer = null; + + assertThat(fires.get()).isEqualTo(0); + } + + @Test + void newTimeout_dispatches_to_supplied_executor() throws Exception { + timer = new ScheduledExecutorTimer("ses-test"); + Thread caller = Thread.currentThread(); + CountDownLatch latch = new CountDownLatch(1); + + timer.newTimeout( + () -> { + assertThat(Thread.currentThread()).isNotSameInstanceAs(caller); + latch.countDown(); + }, + DIRECT, + 5, + TimeUnit.MILLISECONDS); + + assertThat(latch.await(500, TimeUnit.MILLISECONDS)).isTrue(); + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java index 59d9b6e3e702..722ab6a2d4e7 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -72,6 +72,7 @@ import java.time.Instant; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -104,9 +105,7 @@ public class SessionImplTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = - new NettyWheelTimer( - "session-impl-test", com.google.common.util.concurrent.MoreExecutors.directExecutor()); + timer = new ScheduledExecutorTimer("session-impl-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) @@ -826,9 +825,9 @@ private static final class CountingBigtableTimer implements BigtableTimer { } @Override - public Timeout newTimeout(Runnable task, long delay, TimeUnit unit) { + public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { scheduleCount.incrementAndGet(); - Timeout inner = delegate.newTimeout(task, delay, unit); + Timeout inner = delegate.newTimeout(task, executor, delay, unit); return new Timeout() { @Override public boolean cancel() { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 64d9bfbb9ebb..e1ed8a98aa5d 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -54,6 +54,7 @@ import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.base.Suppliers; import com.google.common.truth.Correspondence; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.protobuf.util.Durations; import com.google.rpc.Code; @@ -82,6 +83,7 @@ import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -113,7 +115,7 @@ public class SessionPoolImplTest { Correspondence.transforming(SessionRequest::getOpenSession, "open session"); private ScheduledExecutorService executor; - private NettyWheelTimer testTimer; + private ScheduledExecutorTimer testTimer; @Mock(answer = Answers.RETURNS_DEEP_STUBS) private Metrics metrics; @@ -130,9 +132,7 @@ public class SessionPoolImplTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - testTimer = - new NettyWheelTimer( - "test-timer", com.google.common.util.concurrent.MoreExecutors.directExecutor()); + testTimer = new ScheduledExecutorTimer("test-timer"); fakeService = new FakeSessionService(executor); headerInterceptor = new HeaderInterceptor(); server = @@ -166,7 +166,8 @@ void setUp() throws IOException { CallOptions.DEFAULT, FakeDescriptor.FAKE_SESSION, "fake-pool", - testTimer); + testTimer, + MoreExecutors.directExecutor()); } @AfterEach @@ -250,7 +251,8 @@ void pendingVRpcDrainTest() throws ExecutionException, InterruptedException, Tim CallOptions.DEFAULT, FakeDescriptor.FAKE_SESSION, "fake-pool", - testTimer); + testTimer, + MoreExecutors.directExecutor()); // session ack should be delayed by at least 10ms testSessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); @@ -350,7 +352,8 @@ void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { CallOptions.DEFAULT, FakeDescriptor.FAKE_SESSION, "fake-pool-in-flight", - testTimer); + testTimer, + MoreExecutors.directExecutor()); testSessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); @@ -468,7 +471,8 @@ void testCreateSessionDoesntPropagateDeadline() { CallOptions.DEFAULT, FakeDescriptor.FAKE_SESSION, "fake-pool", - testTimer); + testTimer, + MoreExecutors.directExecutor()); Context.current() .withDeadlineAfter(1, TimeUnit.MINUTES, executor) @@ -529,6 +533,7 @@ void setUp() throws Exception { FakeDescriptor.FAKE_SESSION, "fake-pool", mockTimer, + MoreExecutors.directExecutor(), budget); } @@ -567,7 +572,10 @@ public void test() throws Exception { int waitForReadyMs = 1000; verify(mockTimer, Mockito.timeout(waitForReadyMs)) .newTimeout( - runnableCaptor.capture(), longThat(isRetrySchedule::test), eq(TimeUnit.MILLISECONDS)); + runnableCaptor.capture(), + any(Executor.class), + longThat(isRetrySchedule::test), + eq(TimeUnit.MILLISECONDS)); // we should have received some open requests int requestsBefore = fakeService.getOpenRequestCount().get(); @@ -592,7 +600,10 @@ public void test() throws Exception { // The retry task will try to open new sessions. This will fail and schedule another retry. verify(mockTimer, Mockito.timeout(5000).times(2)) .newTimeout( - any(Runnable.class), longThat(isRetrySchedule::test), eq(TimeUnit.MILLISECONDS)); + any(Runnable.class), + any(Executor.class), + longThat(isRetrySchedule::test), + eq(TimeUnit.MILLISECONDS)); // the retry will exhaust the budget again. we should see double the request compared to // before diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java index 790d997f4295..c95717ccbb4d 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java @@ -26,10 +26,12 @@ import com.google.cloud.bigtable.data.v2.internal.session.Session.SessionState; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolImpl.Watchdog; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeClock; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.Message; import io.grpc.Metadata; import java.time.Duration; import java.time.Instant; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -60,6 +62,7 @@ void setUp() { new Watchdog( new Object(), timer, + MoreExecutors.directExecutor(), interval, sessions, NoopMetrics.NoopDebugTracer.INSTANCE, @@ -70,7 +73,7 @@ void setUp() { // by calling run() directly; the scheduling layer is not under test. private static final class NoOpBigtableTimer implements BigtableTimer { @Override - public Timeout newTimeout(Runnable task, long delay, TimeUnit unit) { + public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { return new Timeout() { @Override public boolean cancel() { From fb43b868de7b165b79d007245eb48b7fdf5b3240 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 30 Jun 2026 00:49:44 +0000 Subject: [PATCH 2/4] refactor: replace ScheduledExecutorTimer with in-process HashedWheelTimer Drops the JDK ScheduledExecutorService backing in favor of a stripped-down in-process hashed-wheel implementation: single tick thread, 512x10ms wheel (~5s rotation), MPSC pending queue, O(1) insert, O(1) cancel via state CAS with lazy reaping on the bucket's next sweep. Same algorithm and tick geometry as the shaded-netty HashedWheelTimer we used pre-phase-7, but with no Netty dependency. BigtableTimer interface and the per-call Executor contract from c29e are preserved unchanged; only the impl swaps. Benchmark (jetstream100, c100, n=3 medians): HashedWheelTimer: 111,557 ops/s p99 1,527us p9999 94ms ScheduledExecutorTimer: 106,868 ops/s p99 1,592us p9999 174ms Wheel recovers the ~4% throughput regression introduced in c29e and roughly halves the p9999 deep tail. --- .../bigtable/data/v2/internal/api/Client.java | 4 +- .../v2/internal/session/HashedWheelTimer.java | 268 ++++++++++++++++++ .../session/ScheduledExecutorTimer.java | 114 -------- .../internal/csm/tracers/VRpcTracerTest.java | 4 +- .../internal/middleware/RetryingVRpcTest.java | 4 +- ...merTest.java => HashedWheelTimerTest.java} | 14 +- .../v2/internal/session/SessionImplTest.java | 2 +- .../internal/session/SessionPoolImplTest.java | 4 +- 8 files changed, 284 insertions(+), 130 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java delete mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java rename java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/{ScheduledExecutorTimerTest.java => HashedWheelTimerTest.java} (91%) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 4ca588403702..8bd9998e86f5 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -34,7 +34,7 @@ import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; -import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; +import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.util.concurrent.ThreadFactoryBuilder; @@ -205,7 +205,7 @@ public Client( this.configManager = configManager; this.backgroundExecutor = bgExecutor; this.userCallbackExecutor = userCallbackExecutor; - this.sessionTimer = new ScheduledExecutorTimer("bigtable-session-timer"); + this.sessionTimer = new HashedWheelTimer("bigtable-session-timer"); defaultCallOptions = CallOptions.DEFAULT; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java new file mode 100644 index 000000000000..13166acf7acf --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java @@ -0,0 +1,268 @@ +/* + * 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.session; + +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import java.util.HashSet; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * {@link BigtableTimer} backed by a single-thread hashed wheel — same algorithm as Netty's + * {@code HashedWheelTimer} but with no Netty dependency. Tick = {@value #TICK_DURATION_MS} ms, + * {@value #TICKS_PER_WHEEL} buckets, so one full rotation is ~5 s. + * + *

Insert is O(1): the new {@link Timeout} is enqueued onto an MPSC pending queue and bucketed + * by the tick thread on the next tick. Cancel is O(1): a single state CAS marks the entry + * cancelled, and the worker unlinks it the next time the bucket is swept (so a cancelled entry + * holds memory for at most one wheel rotation). + * + *

The worker thread only advances the wheel and hops fires to the caller-supplied + * {@link Executor}; user code never runs on the tick thread. + */ +public final class HashedWheelTimer implements BigtableTimer { + private static final Logger LOG = Logger.getLogger(HashedWheelTimer.class.getName()); + + static final long TICK_DURATION_MS = 10; + static final int TICKS_PER_WHEEL = 512; // must be a power of two + private static final long TICK_DURATION_NANOS = TimeUnit.MILLISECONDS.toNanos(TICK_DURATION_MS); + private static final int MASK = TICKS_PER_WHEEL - 1; + // Bound transfers per tick so a sustained insert flood can't starve bucket processing. + private static final int MAX_PENDING_TRANSFERS_PER_TICK = 100_000; + + private final HashedWheelBucket[] wheel = new HashedWheelBucket[TICKS_PER_WHEEL]; + private final Queue pendingTimeouts = new ConcurrentLinkedQueue<>(); + private final Set stopHooks = ConcurrentHashMap.newKeySet(); + + private final Thread worker; + private final long startNanos; + + private volatile boolean stopped = false; + + public HashedWheelTimer(String name) { + for (int i = 0; i < TICKS_PER_WHEEL; i++) { + wheel[i] = new HashedWheelBucket(); + } + this.startNanos = System.nanoTime(); + this.worker = + new ThreadFactoryBuilder() + .setNameFormat(name + "-%d") + .setDaemon(true) + .build() + .newThread(this::workerLoop); + this.worker.start(); + } + + @Override + public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + long delayNanos = Math.max(0L, unit.toNanos(delay)); + long deadlineNanos = (System.nanoTime() - startNanos) + delayNanos; + HashedWheelTimeout timeout = new HashedWheelTimeout(task, executor, deadlineNanos); + pendingTimeouts.add(timeout); + return timeout; + } + + @Override + public Registration onStop(Runnable hook) { + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + stopHooks.add(hook); + return () -> stopHooks.remove(hook); + } + + @Override + public void stop() { + if (stopped) { + return; + } + stopped = true; + Set hooks = new HashSet<>(stopHooks); + stopHooks.clear(); + for (Runnable hook : hooks) { + try { + hook.run(); + } catch (Throwable t) { + LOG.log(Level.WARNING, "stop hook threw; continuing", t); + } + } + worker.interrupt(); + } + + private void workerLoop() { + long tick = 0; + while (!stopped) { + if (waitForNextTick(tick) < 0) { + return; + } + transferPending(tick); + wheel[(int) (tick & MASK)].expireTimeouts(); + tick++; + } + } + + // Returns elapsed nanos since startNanos at wake, or -1 if the timer was stopped. + private long waitForNextTick(long tick) { + long targetNanos = (tick + 1) * TICK_DURATION_NANOS; + while (true) { + long elapsed = System.nanoTime() - startNanos; + long sleepNanos = targetNanos - elapsed; + if (sleepNanos <= 0) { + return elapsed; + } + if (stopped) { + return -1; + } + long sleepMs = (sleepNanos + 999_999L) / 1_000_000L; + try { + Thread.sleep(sleepMs); + } catch (InterruptedException e) { + if (stopped) { + return -1; + } + // Spurious interrupt; recompute and keep sleeping. + } + } + } + + private void transferPending(long tick) { + for (int i = 0; i < MAX_PENDING_TRANSFERS_PER_TICK; i++) { + HashedWheelTimeout t = pendingTimeouts.poll(); + if (t == null) { + return; + } + if (t.isCancelled()) { + continue; + } + long calculated = t.deadlineNanos / TICK_DURATION_NANOS; + t.remainingRounds = Math.max(0L, (calculated - tick) / TICKS_PER_WHEEL); + // If the deadline is already past, place the entry in the current bucket so it fires now. + long ticksTarget = Math.max(calculated, tick); + wheel[(int) (ticksTarget & MASK)].add(t); + } + } + + // Bucket: singly-traversed doubly-linked list, mutated only on the worker thread. + private static final class HashedWheelBucket { + private HashedWheelTimeout head; + private HashedWheelTimeout tail; + + void add(HashedWheelTimeout t) { + t.bucket = this; + if (head == null) { + head = tail = t; + } else { + tail.next = t; + t.prev = tail; + tail = t; + } + } + + void expireTimeouts() { + HashedWheelTimeout t = head; + while (t != null) { + HashedWheelTimeout next = t.next; + if (t.isCancelled()) { + unlink(t); + } else if (t.remainingRounds <= 0) { + unlink(t); + t.expire(); + } else { + t.remainingRounds--; + } + t = next; + } + } + + private void unlink(HashedWheelTimeout t) { + HashedWheelTimeout p = t.prev; + HashedWheelTimeout n = t.next; + if (p != null) { + p.next = n; + } + if (n != null) { + n.prev = p; + } + if (t == head) { + head = n; + } + if (t == tail) { + tail = p; + } + t.prev = null; + t.next = null; + t.bucket = null; + } + } + + private static final class HashedWheelTimeout implements Timeout { + private static final int STATE_INIT = 0; + private static final int STATE_CANCELLED = 1; + private static final int STATE_EXPIRED = 2; + + private static final AtomicIntegerFieldUpdater STATE_UPDATER = + AtomicIntegerFieldUpdater.newUpdater(HashedWheelTimeout.class, "state"); + + private final Runnable task; + private final Executor executor; + final long deadlineNanos; + + volatile int state = STATE_INIT; + // Worker-thread only: + long remainingRounds; + HashedWheelTimeout next; + HashedWheelTimeout prev; + HashedWheelBucket bucket; + + HashedWheelTimeout(Runnable task, Executor executor, long deadlineNanos) { + this.task = task; + this.executor = executor; + this.deadlineNanos = deadlineNanos; + } + + @Override + public boolean cancel() { + return STATE_UPDATER.compareAndSet(this, STATE_INIT, STATE_CANCELLED); + } + + @Override + public boolean isCancelled() { + return state == STATE_CANCELLED; + } + + void expire() { + if (!STATE_UPDATER.compareAndSet(this, STATE_INIT, STATE_EXPIRED)) { + // Lost a race with cancel() between transferPending and this sweep. + return; + } + try { + executor.execute(task); + } catch (Throwable t) { + LOG.log(Level.WARNING, "executor.execute threw; dropping task", t); + } + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java deleted file mode 100644 index 25ce016b73d3..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimer.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * 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.session; - -import com.google.common.util.concurrent.ThreadFactoryBuilder; -import java.util.HashSet; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executor; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; -import java.util.logging.Level; -import java.util.logging.Logger; - -/** - * {@link BigtableTimer} backed by a single-threaded {@link ScheduledExecutorService}. - * - *

Each {@code newTimeout} schedules a wrapper that hops to the caller-supplied {@link Executor}. - * Cancel delegates to {@link ScheduledFuture#cancel(boolean)} with {@code mayInterruptIfRunning = - * false}. - */ -public final class ScheduledExecutorTimer implements BigtableTimer { - private static final Logger LOG = Logger.getLogger(ScheduledExecutorTimer.class.getName()); - - private final ScheduledExecutorService scheduler; - private final Set stopHooks = ConcurrentHashMap.newKeySet(); - - private volatile boolean stopped = false; - - public ScheduledExecutorTimer(String name) { - this.scheduler = - Executors.newSingleThreadScheduledExecutor( - new ThreadFactoryBuilder().setNameFormat(name + "-%d").setDaemon(true).build()); - } - - @Override - public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { - if (stopped) { - throw new IllegalStateException("timer stopped"); - } - ScheduledFuture future = - scheduler.schedule( - () -> { - try { - executor.execute(task); - } catch (Throwable t) { - LOG.log(Level.WARNING, "executor.execute threw; dropping task", t); - } - }, - delay, - unit); - return new FutureTimeout(future); - } - - @Override - public Registration onStop(Runnable hook) { - if (stopped) { - throw new IllegalStateException("timer stopped"); - } - stopHooks.add(hook); - return () -> stopHooks.remove(hook); - } - - @Override - public void stop() { - if (stopped) { - return; - } - stopped = true; - Set hooks = new HashSet<>(stopHooks); - stopHooks.clear(); - for (Runnable hook : hooks) { - try { - hook.run(); - } catch (Throwable t) { - LOG.log(Level.WARNING, "stop hook threw; continuing", t); - } - } - scheduler.shutdownNow(); - } - - private static final class FutureTimeout implements Timeout { - private final ScheduledFuture future; - - FutureTimeout(ScheduledFuture future) { - this.future = future; - } - - @Override - public boolean cancel() { - return future.cancel(false); - } - - @Override - public boolean isCancelled() { - return future.isCancelled(); - } - } -} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java index 1ecc6e88b0d3..2560a93ffab7 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java @@ -49,7 +49,7 @@ import com.google.cloud.bigtable.data.v2.internal.session.Session; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; -import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; +import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; @@ -118,7 +118,7 @@ public class VRpcTracerTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = new ScheduledExecutorTimer("vrpc-tracer-test"); + timer = new HashedWheelTimer("vrpc-tracer-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java index e89f9de108e5..11f073c31b9f 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java @@ -41,7 +41,7 @@ import com.google.cloud.bigtable.data.v2.internal.session.FakeDescriptor; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; -import com.google.cloud.bigtable.data.v2.internal.session.ScheduledExecutorTimer; +import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; @@ -91,7 +91,7 @@ public class RetryingVRpcTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = new ScheduledExecutorTimer("retrying-vrpc-test"); + timer = new HashedWheelTimer("retrying-vrpc-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java similarity index 91% rename from java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java rename to java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java index 189a00903748..619bf3c89fdb 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/ScheduledExecutorTimerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java @@ -31,10 +31,10 @@ import org.junit.jupiter.api.Timeout; @Timeout(value = 10, unit = TimeUnit.SECONDS) -class ScheduledExecutorTimerTest { +class HashedWheelTimerTest { private static final Executor DIRECT = MoreExecutors.directExecutor(); - private ScheduledExecutorTimer timer; + private HashedWheelTimer timer; @AfterEach void tearDown() { @@ -45,7 +45,7 @@ void tearDown() { @Test void schedule_fires_after_delay() throws Exception { - timer = new ScheduledExecutorTimer("ses-test"); + timer = new HashedWheelTimer("wheel-test"); CountDownLatch latch = new CountDownLatch(1); timer.newTimeout(latch::countDown, DIRECT, 10, TimeUnit.MILLISECONDS); @@ -55,7 +55,7 @@ void schedule_fires_after_delay() throws Exception { @Test void cancel_before_fire_prevents_execution() throws Exception { - timer = new ScheduledExecutorTimer("ses-test"); + timer = new HashedWheelTimer("wheel-test"); AtomicInteger fires = new AtomicInteger(); BigtableTimer.Timeout t = @@ -69,7 +69,7 @@ void cancel_before_fire_prevents_execution() throws Exception { @Test void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() { - timer = new ScheduledExecutorTimer("ses-test"); + timer = new HashedWheelTimer("wheel-test"); Thread caller = Thread.currentThread(); List ran = new CopyOnWriteArrayList<>(); @@ -90,7 +90,7 @@ void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() { @Test void unregister_removes_hook() { - timer = new ScheduledExecutorTimer("ses-test"); + timer = new HashedWheelTimer("wheel-test"); AtomicInteger fires = new AtomicInteger(); Registration reg = timer.onStop(fires::incrementAndGet); @@ -104,7 +104,7 @@ void unregister_removes_hook() { @Test void newTimeout_dispatches_to_supplied_executor() throws Exception { - timer = new ScheduledExecutorTimer("ses-test"); + timer = new HashedWheelTimer("wheel-test"); Thread caller = Thread.currentThread(); CountDownLatch latch = new CountDownLatch(1); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java index 722ab6a2d4e7..0b0006135cff 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -105,7 +105,7 @@ public class SessionImplTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - timer = new ScheduledExecutorTimer("session-impl-test"); + timer = new HashedWheelTimer("session-impl-test"); server = FakeServiceBuilder.create(new FakeSessionService(executor)) .intercept(new PeerInfoInterceptor()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index e1ed8a98aa5d..85ebc6937cbe 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -115,7 +115,7 @@ public class SessionPoolImplTest { Correspondence.transforming(SessionRequest::getOpenSession, "open session"); private ScheduledExecutorService executor; - private ScheduledExecutorTimer testTimer; + private HashedWheelTimer testTimer; @Mock(answer = Answers.RETURNS_DEEP_STUBS) private Metrics metrics; @@ -132,7 +132,7 @@ public class SessionPoolImplTest { @BeforeEach void setUp() throws IOException { executor = Executors.newScheduledThreadPool(4); - testTimer = new ScheduledExecutorTimer("test-timer"); + testTimer = new HashedWheelTimer("test-timer"); fakeService = new FakeSessionService(executor); headerInterceptor = new HeaderInterceptor(); server = From 8bd9885a4e52fe952a0917208269c9fe7e044c21 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 30 Jun 2026 16:05:12 +0000 Subject: [PATCH 3/4] chore: apply fmt-maven-plugin --- .../v2/internal/session/HashedWheelTimer.java | 18 +++++++++--------- .../internal/csm/tracers/VRpcTracerTest.java | 2 +- .../internal/middleware/RetryingVRpcTest.java | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java index 13166acf7acf..6da2b2ee4748 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java @@ -28,17 +28,17 @@ import java.util.logging.Logger; /** - * {@link BigtableTimer} backed by a single-thread hashed wheel — same algorithm as Netty's - * {@code HashedWheelTimer} but with no Netty dependency. Tick = {@value #TICK_DURATION_MS} ms, - * {@value #TICKS_PER_WHEEL} buckets, so one full rotation is ~5 s. + * {@link BigtableTimer} backed by a single-thread hashed wheel — same algorithm as Netty's {@code + * HashedWheelTimer} but with no Netty dependency. Tick = {@value #TICK_DURATION_MS} ms, {@value + * #TICKS_PER_WHEEL} buckets, so one full rotation is ~5 s. * - *

Insert is O(1): the new {@link Timeout} is enqueued onto an MPSC pending queue and bucketed - * by the tick thread on the next tick. Cancel is O(1): a single state CAS marks the entry - * cancelled, and the worker unlinks it the next time the bucket is swept (so a cancelled entry - * holds memory for at most one wheel rotation). + *

Insert is O(1): the new {@link Timeout} is enqueued onto an MPSC pending queue and bucketed by + * the tick thread on the next tick. Cancel is O(1): a single state CAS marks the entry cancelled, + * and the worker unlinks it the next time the bucket is swept (so a cancelled entry holds memory + * for at most one wheel rotation). * - *

The worker thread only advances the wheel and hops fires to the caller-supplied - * {@link Executor}; user code never runs on the tick thread. + *

The worker thread only advances the wheel and hops fires to the caller-supplied {@link + * Executor}; user code never runs on the tick thread. */ public final class HashedWheelTimer implements BigtableTimer { private static final Logger LOG = Logger.getLogger(HashedWheelTimer.class.getName()); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java index 2560a93ffab7..54685b907700 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java @@ -46,10 +46,10 @@ import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; import com.google.cloud.bigtable.data.v2.internal.session.FakeDescriptor; +import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.Session; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; -import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java index 11f073c31b9f..0109785ffe98 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java @@ -39,9 +39,9 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; import com.google.cloud.bigtable.data.v2.internal.session.FakeDescriptor; +import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionFactory; import com.google.cloud.bigtable.data.v2.internal.session.SessionImpl; -import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionListener; From 39bbfb4179968a643e876a22a9938a98f1b91fe6 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 30 Jun 2026 17:27:43 +0000 Subject: [PATCH 4/4] fix: close races between HashedWheelTimer stop() and concurrent register-ers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stop() previously did a plain `if (stopped) return; stopped = true` then snapshot+clear of stopHooks. Three races followed: - Double stop() ran every hook twice (both callers passed the early check before either wrote the flag). - onStop() racing with stop() could land an add after the snapshot/clear, leaving the hook stranded in the set with no worker alive to drive it — the caller's Registration claimed the hook would fire on stop but it never would. - newTimeout() racing with stop() left entries in pendingTimeouts that nobody drained. Switch to AtomicBoolean.compareAndSet for the stop transition. After the CAS, drain stopHooks AND pendingTimeouts in a loop until both observe empty, so any add that started before the CAS but landed after the first drain pass gets caught on the next pass. onStop now invokes the hook synchronously when stop() already happened (ListenableFuture addListener semantics). newTimeout returns a pre-cancelled singleton in the same case so callers see `isCancelled() == true` and skip cleanup paths. The register-then-recheck pattern in each handles the stop()-just-CAS'd-but-not-yet-drained window: the late add either lands in the queue/set (where stop()'s drain loop catches it) or sees the flag and self-removes/self-fires. Drops Scheduled.onStart's try/catch around the timer calls — neither onStop nor newTimeout can throw on the stopped path anymore. --- .../v2/internal/middleware/RetryingVRpc.java | 40 ++++------ .../v2/internal/session/HashedWheelTimer.java | 80 +++++++++++++++---- .../session/HashedWheelTimerTest.java | 39 +++++++-- 3 files changed, 112 insertions(+), 47 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java index 434577bc138b..206ff052bbdc 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java @@ -326,31 +326,21 @@ class Scheduled extends State { @Override public void onStart() { - try { - stopHook = timer.onStop(this::onTimerStopping); - // Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at - // the moment the body runs on the op executor. - future = - timer.newTimeout( - () -> - grpcContext - .wrap(() -> otelContext.wrap(() -> onStateChange(new Idle())).run()) - .run(), - context.getExecutor(), - Durations.toMillis(retryDelay), - TimeUnit.MILLISECONDS); - } catch (IllegalStateException e) { - // Timer was stopped between Active.onClose deciding to retry and this task running on the - // op executor. Race window is narrow (post-drain shutdown), but cover it cleanly so the - // op-executor uncaught handler does not have to. onExit will release the stopHook. - onStateChange( - new Done( - VRpcResult.createRejectedError( - Status.CANCELLED - .withDescription( - "Executor shutting down, can't schedule operation for retry.") - .withCause(e)))); - } + // If the timer is already stopped, onStop fires the hook synchronously here — which + // enqueues an onTimerStopping body on the op executor — and newTimeout returns a + // pre-cancelled Timeout. The queued onTimerStopping then drives us to Done(CANCELLED). + stopHook = timer.onStop(this::onTimerStopping); + // Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at + // the moment the body runs on the op executor. + future = + timer.newTimeout( + () -> + grpcContext + .wrap(() -> otelContext.wrap(() -> onStateChange(new Idle())).run()) + .run(), + context.getExecutor(), + Durations.toMillis(retryDelay), + TimeUnit.MILLISECONDS); } // Invoked from BigtableTimer.stop on the close thread. Trampoline back to the op executor so diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java index 6da2b2ee4748..33edb1659a4b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java @@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import java.util.logging.Level; import java.util.logging.Logger; @@ -50,6 +51,21 @@ public final class HashedWheelTimer implements BigtableTimer { // Bound transfers per tick so a sustained insert flood can't starve bucket processing. private static final int MAX_PENDING_TRANSFERS_PER_TICK = 100_000; + // Returned from newTimeout when the timer is already stopped. Pre-cancelled so callers that + // poll isCancelled() see "done" and skip cleanup paths; cancel() is a no-op. + private static final Timeout DEAD_TIMEOUT = + new Timeout() { + @Override + public boolean cancel() { + return false; + } + + @Override + public boolean isCancelled() { + return true; + } + }; + private final HashedWheelBucket[] wheel = new HashedWheelBucket[TICKS_PER_WHEEL]; private final Queue pendingTimeouts = new ConcurrentLinkedQueue<>(); private final Set stopHooks = ConcurrentHashMap.newKeySet(); @@ -57,7 +73,10 @@ public final class HashedWheelTimer implements BigtableTimer { private final Thread worker; private final long startNanos; - private volatile boolean stopped = false; + // Single-shot transition from false to true; CAS in stop() pairs with .get() in newTimeout/ + // onStop so racing register-ers either land before stop()'s drain (and are caught by it) or + // observe true and self-resolve. + private final AtomicBoolean stopped = new AtomicBoolean(false); public HashedWheelTimer(String name) { for (int i = 0; i < TICKS_PER_WHEEL; i++) { @@ -75,38 +94,65 @@ public HashedWheelTimer(String name) { @Override public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { - if (stopped) { - throw new IllegalStateException("timer stopped"); + if (stopped.get()) { + return DEAD_TIMEOUT; } long delayNanos = Math.max(0L, unit.toNanos(delay)); long deadlineNanos = (System.nanoTime() - startNanos) + delayNanos; HashedWheelTimeout timeout = new HashedWheelTimeout(task, executor, deadlineNanos); pendingTimeouts.add(timeout); + // If stop() raced after our get(), its drain loop will catch this entry and cancel it. return timeout; } @Override public Registration onStop(Runnable hook) { - if (stopped) { - throw new IllegalStateException("timer stopped"); + if (stopped.get()) { + hook.run(); + return NO_OP_REGISTRATION; } stopHooks.add(hook); + if (stopped.get()) { + // Raced with stop(). Either its drain already consumed and invoked our hook (set.remove + // returns false → no-op), or our add landed after its drain (set.remove returns true → + // we own it and must fire it ourselves). Either way the caller sees the hook run. + if (stopHooks.remove(hook)) { + hook.run(); + } + return NO_OP_REGISTRATION; + } return () -> stopHooks.remove(hook); } + private static final Registration NO_OP_REGISTRATION = () -> {}; + @Override public void stop() { - if (stopped) { + if (!stopped.compareAndSet(false, true)) { return; } - stopped = true; - Set hooks = new HashSet<>(stopHooks); - stopHooks.clear(); - for (Runnable hook : hooks) { - try { - hook.run(); - } catch (Throwable t) { - LOG.log(Level.WARNING, "stop hook threw; continuing", t); + // Drain both queues in a loop. In-flight register-ers that read stopped==false before our + // CAS will eventually finish their add and land in stopHooks/pendingTimeouts; the loop + // keeps draining until no new entries appear. Readers that arrive after our CAS observe + // stopped==true and skip the add entirely. + while (true) { + Set hooks = new HashSet<>(stopHooks); + stopHooks.removeAll(hooks); + for (Runnable hook : hooks) { + try { + hook.run(); + } catch (Throwable t) { + LOG.log(Level.WARNING, "stop hook threw; continuing", t); + } + } + int drained = 0; + HashedWheelTimeout t; + while ((t = pendingTimeouts.poll()) != null) { + t.cancel(); + drained++; + } + if (hooks.isEmpty() && drained == 0) { + break; } } worker.interrupt(); @@ -114,7 +160,7 @@ public void stop() { private void workerLoop() { long tick = 0; - while (!stopped) { + while (!stopped.get()) { if (waitForNextTick(tick) < 0) { return; } @@ -133,14 +179,14 @@ private long waitForNextTick(long tick) { if (sleepNanos <= 0) { return elapsed; } - if (stopped) { + if (stopped.get()) { return -1; } long sleepMs = (sleepNanos + 999_999L) / 1_000_000L; try { Thread.sleep(sleepMs); } catch (InterruptedException e) { - if (stopped) { + if (stopped.get()) { return -1; } // Spurious interrupt; recompute and keep sleeping. diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java index 619bf3c89fdb..b092c3891ba7 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java @@ -16,7 +16,6 @@ package com.google.cloud.bigtable.data.v2.internal.session; import static com.google.common.truth.Truth.assertThat; -import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer.Registration; import com.google.common.util.concurrent.MoreExecutors; @@ -68,7 +67,7 @@ void cancel_before_fire_prevents_execution() throws Exception { } @Test - void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() { + void stop_runs_all_hooks_on_caller_thread_then_newTimeout_returns_dead() { timer = new HashedWheelTimer("wheel-test"); Thread caller = Thread.currentThread(); List ran = new CopyOnWriteArrayList<>(); @@ -82,9 +81,39 @@ void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() { for (Thread t : ran) { assertThat(t).isSameInstanceAs(caller); } - assertThrows( - IllegalStateException.class, - () -> timer.newTimeout(() -> {}, DIRECT, 1, TimeUnit.MILLISECONDS)); + // After stop, newTimeout returns a pre-cancelled Timeout that never fires. + BigtableTimer.Timeout t = timer.newTimeout(() -> {}, DIRECT, 1, TimeUnit.MILLISECONDS); + assertThat(t.isCancelled()).isTrue(); + timer = null; + } + + @Test + void onStop_after_stop_fires_hook_synchronously() { + timer = new HashedWheelTimer("wheel-test"); + timer.stop(); + Thread caller = Thread.currentThread(); + List ran = new CopyOnWriteArrayList<>(); + + Registration reg = timer.onStop(() -> ran.add(Thread.currentThread())); + + assertThat(ran).hasSize(1); + assertThat(ran.get(0)).isSameInstanceAs(caller); + // Registration is still safe to call; should be a no-op. + reg.unregister(); + timer = null; + } + + @Test + void double_stop_does_not_re_run_hooks() { + timer = new HashedWheelTimer("wheel-test"); + AtomicInteger fires = new AtomicInteger(); + + timer.onStop(fires::incrementAndGet); + + timer.stop(); + timer.stop(); + + assertThat(fires.get()).isEqualTo(1); timer = null; }