Skip to content

Commit 39bbfb4

Browse files
fix: close races between HashedWheelTimer stop() and concurrent register-ers
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.
1 parent 8bd9885 commit 39bbfb4

3 files changed

Lines changed: 112 additions & 47 deletions

File tree

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -326,31 +326,21 @@ class Scheduled extends State {
326326

327327
@Override
328328
public void onStart() {
329-
try {
330-
stopHook = timer.onStop(this::onTimerStopping);
331-
// Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at
332-
// the moment the body runs on the op executor.
333-
future =
334-
timer.newTimeout(
335-
() ->
336-
grpcContext
337-
.wrap(() -> otelContext.wrap(() -> onStateChange(new Idle())).run())
338-
.run(),
339-
context.getExecutor(),
340-
Durations.toMillis(retryDelay),
341-
TimeUnit.MILLISECONDS);
342-
} catch (IllegalStateException e) {
343-
// Timer was stopped between Active.onClose deciding to retry and this task running on the
344-
// op executor. Race window is narrow (post-drain shutdown), but cover it cleanly so the
345-
// op-executor uncaught handler does not have to. onExit will release the stopHook.
346-
onStateChange(
347-
new Done(
348-
VRpcResult.createRejectedError(
349-
Status.CANCELLED
350-
.withDescription(
351-
"Executor shutting down, can't schedule operation for retry.")
352-
.withCause(e))));
353-
}
329+
// If the timer is already stopped, onStop fires the hook synchronously here — which
330+
// enqueues an onTimerStopping body on the op executor — and newTimeout returns a
331+
// pre-cancelled Timeout. The queued onTimerStopping then drives us to Done(CANCELLED).
332+
stopHook = timer.onStop(this::onTimerStopping);
333+
// Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at
334+
// the moment the body runs on the op executor.
335+
future =
336+
timer.newTimeout(
337+
() ->
338+
grpcContext
339+
.wrap(() -> otelContext.wrap(() -> onStateChange(new Idle())).run())
340+
.run(),
341+
context.getExecutor(),
342+
Durations.toMillis(retryDelay),
343+
TimeUnit.MILLISECONDS);
354344
}
355345

356346
// Invoked from BigtableTimer.stop on the close thread. Trampoline back to the op executor so

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimer.java‎

Lines changed: 63 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.concurrent.ConcurrentLinkedQueue;
2424
import java.util.concurrent.Executor;
2525
import java.util.concurrent.TimeUnit;
26+
import java.util.concurrent.atomic.AtomicBoolean;
2627
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
2728
import java.util.logging.Level;
2829
import java.util.logging.Logger;
@@ -50,14 +51,32 @@ public final class HashedWheelTimer implements BigtableTimer {
5051
// Bound transfers per tick so a sustained insert flood can't starve bucket processing.
5152
private static final int MAX_PENDING_TRANSFERS_PER_TICK = 100_000;
5253

54+
// Returned from newTimeout when the timer is already stopped. Pre-cancelled so callers that
55+
// poll isCancelled() see "done" and skip cleanup paths; cancel() is a no-op.
56+
private static final Timeout DEAD_TIMEOUT =
57+
new Timeout() {
58+
@Override
59+
public boolean cancel() {
60+
return false;
61+
}
62+
63+
@Override
64+
public boolean isCancelled() {
65+
return true;
66+
}
67+
};
68+
5369
private final HashedWheelBucket[] wheel = new HashedWheelBucket[TICKS_PER_WHEEL];
5470
private final Queue<HashedWheelTimeout> pendingTimeouts = new ConcurrentLinkedQueue<>();
5571
private final Set<Runnable> stopHooks = ConcurrentHashMap.newKeySet();
5672

5773
private final Thread worker;
5874
private final long startNanos;
5975

60-
private volatile boolean stopped = false;
76+
// Single-shot transition from false to true; CAS in stop() pairs with .get() in newTimeout/
77+
// onStop so racing register-ers either land before stop()'s drain (and are caught by it) or
78+
// observe true and self-resolve.
79+
private final AtomicBoolean stopped = new AtomicBoolean(false);
6180

6281
public HashedWheelTimer(String name) {
6382
for (int i = 0; i < TICKS_PER_WHEEL; i++) {
@@ -75,46 +94,73 @@ public HashedWheelTimer(String name) {
7594

7695
@Override
7796
public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) {
78-
if (stopped) {
79-
throw new IllegalStateException("timer stopped");
97+
if (stopped.get()) {
98+
return DEAD_TIMEOUT;
8099
}
81100
long delayNanos = Math.max(0L, unit.toNanos(delay));
82101
long deadlineNanos = (System.nanoTime() - startNanos) + delayNanos;
83102
HashedWheelTimeout timeout = new HashedWheelTimeout(task, executor, deadlineNanos);
84103
pendingTimeouts.add(timeout);
104+
// If stop() raced after our get(), its drain loop will catch this entry and cancel it.
85105
return timeout;
86106
}
87107

88108
@Override
89109
public Registration onStop(Runnable hook) {
90-
if (stopped) {
91-
throw new IllegalStateException("timer stopped");
110+
if (stopped.get()) {
111+
hook.run();
112+
return NO_OP_REGISTRATION;
92113
}
93114
stopHooks.add(hook);
115+
if (stopped.get()) {
116+
// Raced with stop(). Either its drain already consumed and invoked our hook (set.remove
117+
// returns false → no-op), or our add landed after its drain (set.remove returns true →
118+
// we own it and must fire it ourselves). Either way the caller sees the hook run.
119+
if (stopHooks.remove(hook)) {
120+
hook.run();
121+
}
122+
return NO_OP_REGISTRATION;
123+
}
94124
return () -> stopHooks.remove(hook);
95125
}
96126

127+
private static final Registration NO_OP_REGISTRATION = () -> {};
128+
97129
@Override
98130
public void stop() {
99-
if (stopped) {
131+
if (!stopped.compareAndSet(false, true)) {
100132
return;
101133
}
102-
stopped = true;
103-
Set<Runnable> hooks = new HashSet<>(stopHooks);
104-
stopHooks.clear();
105-
for (Runnable hook : hooks) {
106-
try {
107-
hook.run();
108-
} catch (Throwable t) {
109-
LOG.log(Level.WARNING, "stop hook threw; continuing", t);
134+
// Drain both queues in a loop. In-flight register-ers that read stopped==false before our
135+
// CAS will eventually finish their add and land in stopHooks/pendingTimeouts; the loop
136+
// keeps draining until no new entries appear. Readers that arrive after our CAS observe
137+
// stopped==true and skip the add entirely.
138+
while (true) {
139+
Set<Runnable> hooks = new HashSet<>(stopHooks);
140+
stopHooks.removeAll(hooks);
141+
for (Runnable hook : hooks) {
142+
try {
143+
hook.run();
144+
} catch (Throwable t) {
145+
LOG.log(Level.WARNING, "stop hook threw; continuing", t);
146+
}
147+
}
148+
int drained = 0;
149+
HashedWheelTimeout t;
150+
while ((t = pendingTimeouts.poll()) != null) {
151+
t.cancel();
152+
drained++;
153+
}
154+
if (hooks.isEmpty() && drained == 0) {
155+
break;
110156
}
111157
}
112158
worker.interrupt();
113159
}
114160

115161
private void workerLoop() {
116162
long tick = 0;
117-
while (!stopped) {
163+
while (!stopped.get()) {
118164
if (waitForNextTick(tick) < 0) {
119165
return;
120166
}
@@ -133,14 +179,14 @@ private long waitForNextTick(long tick) {
133179
if (sleepNanos <= 0) {
134180
return elapsed;
135181
}
136-
if (stopped) {
182+
if (stopped.get()) {
137183
return -1;
138184
}
139185
long sleepMs = (sleepNanos + 999_999L) / 1_000_000L;
140186
try {
141187
Thread.sleep(sleepMs);
142188
} catch (InterruptedException e) {
143-
if (stopped) {
189+
if (stopped.get()) {
144190
return -1;
145191
}
146192
// Spurious interrupt; recompute and keep sleeping.

‎java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/HashedWheelTimerTest.java‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
package com.google.cloud.bigtable.data.v2.internal.session;
1717

1818
import static com.google.common.truth.Truth.assertThat;
19-
import static org.junit.jupiter.api.Assertions.assertThrows;
2019

2120
import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer.Registration;
2221
import com.google.common.util.concurrent.MoreExecutors;
@@ -68,7 +67,7 @@ void cancel_before_fire_prevents_execution() throws Exception {
6867
}
6968

7069
@Test
71-
void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() {
70+
void stop_runs_all_hooks_on_caller_thread_then_newTimeout_returns_dead() {
7271
timer = new HashedWheelTimer("wheel-test");
7372
Thread caller = Thread.currentThread();
7473
List<Thread> ran = new CopyOnWriteArrayList<>();
@@ -82,9 +81,39 @@ void stop_runs_all_hooks_on_caller_thread_and_rejects_inserts() {
8281
for (Thread t : ran) {
8382
assertThat(t).isSameInstanceAs(caller);
8483
}
85-
assertThrows(
86-
IllegalStateException.class,
87-
() -> timer.newTimeout(() -> {}, DIRECT, 1, TimeUnit.MILLISECONDS));
84+
// After stop, newTimeout returns a pre-cancelled Timeout that never fires.
85+
BigtableTimer.Timeout t = timer.newTimeout(() -> {}, DIRECT, 1, TimeUnit.MILLISECONDS);
86+
assertThat(t.isCancelled()).isTrue();
87+
timer = null;
88+
}
89+
90+
@Test
91+
void onStop_after_stop_fires_hook_synchronously() {
92+
timer = new HashedWheelTimer("wheel-test");
93+
timer.stop();
94+
Thread caller = Thread.currentThread();
95+
List<Thread> ran = new CopyOnWriteArrayList<>();
96+
97+
Registration reg = timer.onStop(() -> ran.add(Thread.currentThread()));
98+
99+
assertThat(ran).hasSize(1);
100+
assertThat(ran.get(0)).isSameInstanceAs(caller);
101+
// Registration is still safe to call; should be a no-op.
102+
reg.unregister();
103+
timer = null;
104+
}
105+
106+
@Test
107+
void double_stop_does_not_re_run_hooks() {
108+
timer = new HashedWheelTimer("wheel-test");
109+
AtomicInteger fires = new AtomicInteger();
110+
111+
timer.onStop(fires::incrementAndGet);
112+
113+
timer.stop();
114+
timer.stop();
115+
116+
assertThat(fires.get()).isEqualTo(1);
88117
timer = null;
89118
}
90119

0 commit comments

Comments
 (0)