This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f353f12b543 [Dataflow Streaming] Mark worker as unhealthy in presence
of stuck commits (#39666)
f353f12b543 is described below
commit f353f12b543e50de2730ec15ca7964889252e4ce
Author: Arun Pandian <[email protected]>
AuthorDate: Tue Aug 11 04:23:50 2026 -0700
[Dataflow Streaming] Mark worker as unhealthy in presence of stuck commits
(#39666)
* [Dataflow Streaming] Mark worker as unhealthy in presence of stuck commits
Getting rid of the commit invalidation logic simplifies work lifecycle.
When worker is marked as unhealthy backend will restart the worker.
* address commits
* fix test
---
.../options/DataflowStreamingPipelineOptions.java | 2 +-
.../dataflow/worker/StreamingDataflowWorker.java | 34 ++++++++---
.../dataflow/worker/streaming/ActiveWorkState.java | 46 ++++-----------
.../worker/streaming/ComputationState.java | 5 +-
.../windmill/work/refresh/ActiveWorkRefresher.java | 21 +------
.../worker/StreamingDataflowWorkerTest.java | 38 ++++--------
.../worker/streaming/ActiveWorkStateTest.java | 42 +++++--------
.../work/refresh/ActiveWorkRefresherTest.java | 69 ----------------------
8 files changed, 74 insertions(+), 183 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
index cca261fff62..58521d9b126 100644
---
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
+++
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
@@ -190,7 +190,7 @@ public interface DataflowStreamingPipelineOptions extends
PipelineOptions {
void setMaxStackTraceDepthToReport(int value);
- @Description("Necessary duration for a commit to be considered stuck and
invalidated.")
+ @Description("Necessary duration for a commit to be considered stuck and
mark worker unhealthy.")
@Default.Integer(60 * 60 * 1000)
int getStuckCommitDurationMillis();
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index 2339430464c..c10e3c0c488 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
@@ -35,6 +35,7 @@ import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
@@ -311,9 +312,6 @@ public final class StreamingDataflowWorker {
new ActiveWorkRefresher(
clock,
options.getActiveWorkRefreshPeriodMillis(),
- options.isEnableStreamingEngine()
- ? Math.max(options.getStuckCommitDurationMillis(), 0)
- : 0,
computationStateCache::getAllPresentComputations,
sampler,
activeWorkRefreshExecutorFn,
@@ -321,7 +319,10 @@ public final class StreamingDataflowWorker {
this.statusPages.set(
createStatusPageBuilder(
- this.options, this.windmillStreamFactory,
this.memoryMonitor.memoryMonitor())
+ this.options,
+ this.windmillStreamFactory,
+ this.memoryMonitor.memoryMonitor(),
+ this::isHealthy)
.setClock(this.clock)
.setClientId(this.clientId)
.setIsRunning(this.running)
@@ -573,7 +574,10 @@ public final class StreamingDataflowWorker {
}
this.statusPages.set(
createStatusPageBuilder(
- this.options, this.windmillStreamFactory,
this.memoryMonitor.memoryMonitor())
+ this.options,
+ this.windmillStreamFactory,
+ this.memoryMonitor.memoryMonitor(),
+ this::isHealthy)
.setClock(this.clock)
.setClientId(this.clientId)
.setIsRunning(this.running)
@@ -590,12 +594,28 @@ public final class StreamingDataflowWorker {
LOG.info("Started new StreamingWorkerStatusPages instance.");
}
+ @VisibleForTesting
+ boolean isHealthy() {
+ int stuckCommitDurationMillis =
+ options.isEnableStreamingEngine() ?
Math.max(options.getStuckCommitDurationMillis(), 0) : 0;
+ if (stuckCommitDurationMillis > 0) {
+ Instant stuckCommitDeadline =
clock.get().minus(Duration.millis(stuckCommitDurationMillis));
+ for (ComputationState computationState :
computationStateCache.getAllPresentComputations()) {
+ if (computationState.hasStuckCommits(stuckCommitDeadline)) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
private static StreamingWorkerStatusPages.Builder createStatusPageBuilder(
DataflowWorkerHarnessOptions options,
GrpcWindmillStreamFactory windmillStreamFactory,
- MemoryMonitor memoryMonitor) {
+ MemoryMonitor memoryMonitor,
+ BooleanSupplier healthyIndicator) {
WorkerStatusPages workerStatusPages =
- WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor);
+ WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor,
healthyIndicator);
StreamingWorkerStatusPages.Builder streamingStatusPages =
StreamingWorkerStatusPages.builder().setStatusPages(workerStatusPages);
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java
index de408258129..1e9c8629b9f 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java
@@ -27,7 +27,6 @@ import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Queue;
-import java.util.function.BiConsumer;
import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
@@ -39,7 +38,6 @@ import org.apache.beam.sdk.annotations.Internal;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
@@ -266,46 +264,28 @@ public final class ActiveWorkState {
return nextWork;
}
- /**
- * Invalidates all {@link Work} that is in the {@link Work.State#COMMITTING}
state which started
- * before the stuckCommitDeadline.
- */
- synchronized void invalidateStuckCommits(
- Instant stuckCommitDeadline, BiConsumer<ShardedKey, WorkId>
shardedKeyAndWorkTokenConsumer) {
- for (Entry<ShardedKey, WorkId> shardedKeyAndWorkId :
- getStuckCommitsAt(stuckCommitDeadline).entrySet()) {
- ShardedKey shardedKey = shardedKeyAndWorkId.getKey();
- WorkId workId = shardedKeyAndWorkId.getValue();
- computationStateCache.invalidate(shardedKey.key(),
shardedKey.shardingKey());
- shardedKeyAndWorkTokenConsumer.accept(shardedKey, workId);
- }
- }
-
private static @Nullable ExecutableWork firstValue(Map<WorkId,
ExecutableWork> map) {
Iterator<Entry<WorkId, ExecutableWork>> iterator =
map.entrySet().iterator();
return iterator.hasNext() ? iterator.next().getValue() : null;
}
- private synchronized ImmutableMap<ShardedKey, WorkId> getStuckCommitsAt(
- Instant stuckCommitDeadline) {
- // Determine the stuck commit keys but complete them outside the loop
iterating over
- // activeWork as completeWork may delete the entry from activeWork.
- ImmutableMap.Builder<ShardedKey, WorkId> stuckCommits =
ImmutableMap.builder();
+ /**
+ * Returns true if there is any {@link Work} in the {@link
Work.State#COMMITTING} state which
+ * started before the stuckCommitDeadline.
+ */
+ synchronized boolean hasStuckCommits(Instant stuckCommitDeadline) {
for (Entry<Long, LinkedHashMap<WorkId, ExecutableWork>> entry :
activeWork.entrySet()) {
@Nullable ExecutableWork executableWork = firstValue(entry.getValue());
- if (executableWork != null) {
- Work work = executableWork.work();
- if (work.isStuckCommittingAt(stuckCommitDeadline)) {
- LOG.error(
- "Detected key {} stuck in COMMITTING state since {}, completing
it with error.",
- work.getShardedKey(),
- work.getStateStartTime());
- stuckCommits.put(work.getShardedKey(), work.id());
- }
+ if (executableWork != null
+ && executableWork.work().isStuckCommittingAt(stuckCommitDeadline)) {
+ LOG.warn(
+ "Detected key {} stuck in COMMITTING state since {}",
+ executableWork.work().getShardedKey(),
+ executableWork.work().getStateStartTime());
+ return true;
}
}
-
- return stuckCommits.build();
+ return false;
}
/**
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
index 5e850d4312e..e2333e23a71 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
@@ -142,9 +142,8 @@ public class ComputationState {
}
}
- public void invalidateStuckCommits(Instant stuckCommitDeadline) {
- activeWorkState.invalidateStuckCommits(
- stuckCommitDeadline, this::completeWorkAndScheduleNextWorkForKey);
+ public boolean hasStuckCommits(Instant stuckCommitDeadline) {
+ return activeWorkState.hasStuckCommits(stuckCommitDeadline);
}
private void execute(ExecutableWork executableWork) {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java
index 793a1aea95e..d030df7f252 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java
@@ -44,8 +44,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
- * Asynchronously GetData requests to Streaming Engine for all sufficiently
old active work and
- * invalidates stuck commits.
+ * Asynchronously GetData requests to Streaming Engine for all sufficiently
old active work.
*
* <p>This informs Windmill that processing is ongoing and the work should not
be retried. The age
* threshold is determined by {@link #activeWorkRefreshPeriodMillis}
@@ -61,7 +60,6 @@ public final class ActiveWorkRefresher {
private final int activeWorkRefreshPeriodMillis;
private final Supplier<Collection<ComputationState>> computations;
private final DataflowExecutionStateSampler sampler;
- private final int stuckCommitDurationMillis;
private final HeartbeatTracker heartbeatTracker;
private final ScheduledExecutorService activeWorkRefreshExecutor;
private final ExecutorService fanOutActiveWorkRefreshExecutor;
@@ -69,14 +67,12 @@ public final class ActiveWorkRefresher {
public ActiveWorkRefresher(
Supplier<Instant> clock,
int activeWorkRefreshPeriodMillis,
- int stuckCommitDurationMillis,
Supplier<Collection<ComputationState>> computations,
DataflowExecutionStateSampler sampler,
ScheduledExecutorService activeWorkRefreshExecutor,
HeartbeatTracker heartbeatTracker) {
this.clock = clock;
this.activeWorkRefreshPeriodMillis = activeWorkRefreshPeriodMillis;
- this.stuckCommitDurationMillis = stuckCommitDurationMillis;
this.computations = computations;
this.sampler = sampler;
this.activeWorkRefreshExecutor = activeWorkRefreshExecutor;
@@ -101,16 +97,10 @@ public final class ActiveWorkRefresher {
activeWorkRefreshPeriodMillis,
TimeUnit.MILLISECONDS);
}
-
- if (stuckCommitDurationMillis > 0) {
- int periodMillis = Math.max(stuckCommitDurationMillis / 10, 100);
- activeWorkRefreshExecutor.scheduleWithFixedDelay(
- this::invalidateStuckCommits, periodMillis, periodMillis,
TimeUnit.MILLISECONDS);
- }
}
public void stop() {
- if (activeWorkRefreshPeriodMillis > 0 || stuckCommitDurationMillis > 0) {
+ if (activeWorkRefreshPeriodMillis > 0) {
activeWorkRefreshExecutor.shutdown();
try {
activeWorkRefreshExecutor.awaitTermination(300, TimeUnit.SECONDS);
@@ -120,13 +110,6 @@ public final class ActiveWorkRefresher {
}
}
- private void invalidateStuckCommits() {
- Instant stuckCommitDeadline =
clock.get().minus(Duration.millis(stuckCommitDurationMillis));
- for (ComputationState computationState : computations.get()) {
- computationState.invalidateStuckCommits(stuckCommitDeadline);
- }
- }
-
private void refreshActiveWork() {
Instant refreshDeadline =
clock.get().minus(Duration.millis(activeWorkRefreshPeriodMillis));
Map<HeartbeatSender, Heartbeats> heartbeatsBySender =
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
index 9ed705550bc..2d35da51a79 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
@@ -74,7 +74,6 @@ import java.util.Optional;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.concurrent.Callable;
-import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
@@ -118,7 +117,6 @@ import
org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames;
import
org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
-import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataResponse;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationHeartbeatRequest;
@@ -4450,10 +4448,9 @@ public class StreamingDataflowWorkerTest {
}
@Test
- public void testStuckCommit() throws Exception {
+ public void testStuckCommitMarksWorkerUnhealthy() throws Exception {
if (!streamingEngine) {
- // Stuck commits have only been observed with streaming engine and thus
recovery from them is
- // not implemented for non-streaming engine.
+ // Stuck commits have only been observed with streaming engine.
return;
}
@@ -4462,40 +4459,31 @@ public class StreamingDataflowWorkerTest {
makeSourceInstruction(StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
+ FakeClock clock = new FakeClock();
StreamingDataflowWorker worker =
makeWorker(
defaultWorkerParams("--stuckCommitDurationMillis=2000")
.setInstructions(instructions)
+ .setClock(clock)
.publishCounters()
.build());
worker.start();
+ assertTrue(worker.isHealthy());
+
// Prevent commit callbacks from being called to simulate a stuck commit.
server.setDropStreamingCommits(true);
- // Add some work for key 1.
+ // Add work to trigger a commit that will get stuck.
server
.whenGetWorkCalled()
- .thenReturn(makeInput(10, TimeUnit.MILLISECONDS.toMicros(2),
DEFAULT_KEY_STRING, 1))
- .thenReturn(makeInput(15, TimeUnit.MILLISECONDS.toMicros(3),
DEFAULT_KEY_STRING, 5));
- ConcurrentHashMap<Long, Consumer<CommitStatus>> droppedCommits =
- server.waitForDroppedCommits(2);
+ .thenReturn(makeInput(10, TimeUnit.MILLISECONDS.toMicros(2),
DEFAULT_KEY_STRING, 1));
+ server.waitForDroppedCommits(1);
+
+ clock.sleep(Duration.millis(3000));
+
+ assertFalse(worker.isHealthy());
server.setDropStreamingCommits(false);
- // Enqueue another work item for key 1.
- server
- .whenGetWorkCalled()
- .thenReturn(makeInput(1, TimeUnit.MILLISECONDS.toMicros(1),
DEFAULT_KEY_STRING, 1));
- // Ensure that this work item processes.
- Map<Long, Windmill.WorkItemCommitRequest> result =
server.waitForAndGetCommits(1);
- // Now ensure that nothing happens if a dropped commit actually completes.
- droppedCommits.values().iterator().next().accept(CommitStatus.OK);
worker.stop();
-
- assertTrue(result.containsKey(1L));
- assertEquals(
- makeExpectedOutput(
- 1, TimeUnit.MILLISECONDS.toMicros(1), DEFAULT_KEY_STRING, 1,
DEFAULT_KEY_STRING)
- .build(),
- removeDynamicFields(result.get(1L)));
}
@Test
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
index aa0eae0d159..1826b7817fd 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
@@ -25,7 +25,6 @@ import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
import java.util.Arrays;
import java.util.Collections;
@@ -36,6 +35,7 @@ import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Supplier;
import
org.apache.beam.runners.dataflow.worker.streaming.ActiveWorkState.ActivateWorkResult;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import
org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient;
@@ -44,6 +44,7 @@ import
org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudge
import
org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
@@ -66,7 +67,7 @@ public class ActiveWorkStateTest {
return ShardedKey.create(ByteString.copyFromUtf8(str), shardKey);
}
- private static ExecutableWork createWork(Windmill.WorkItem workItem) {
+ private static ExecutableWork createWork(Windmill.WorkItem workItem,
Supplier<Instant> clock) {
return ExecutableWork.create(
Work.create(
workItem,
@@ -74,22 +75,13 @@ public class ActiveWorkStateTest {
Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(),
createWorkProcessingContext(),
false,
- Instant::now,
+ clock,
ImmutableList.of()),
(work, handle) -> {});
}
- private static ExecutableWork expiredWork(Windmill.WorkItem workItem) {
- return ExecutableWork.create(
- Work.create(
- workItem,
- workItem.getSerializedSize(),
- Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(),
- createWorkProcessingContext(),
- false,
- () -> Instant.EPOCH,
- ImmutableList.of()),
- (work, handle) -> {});
+ private static ExecutableWork createWork(Windmill.WorkItem workItem) {
+ return createWork(workItem, Instant::now);
}
private static Work.ProcessingContext createWorkProcessingContext() {
@@ -319,25 +311,23 @@ public class ActiveWorkStateTest {
}
@Test
- public void testInvalidateStuckCommits() {
- Map<ShardedKey, WorkId> invalidatedCommits = new HashMap<>();
+ public void testHasStuckCommits() {
ShardedKey shardedKey1 = shardedKey("someKey", 1L);
ShardedKey shardedKey2 = shardedKey("anotherKey", 2L);
- ExecutableWork stuckWork1 = expiredWork(createWorkItem(1L, 1L,
shardedKey1));
- stuckWork1.work().setState(Work.State.COMMITTING);
- ExecutableWork stuckWork2 = expiredWork(createWorkItem(2L, 1L,
shardedKey2));
- stuckWork2.work().setState(Work.State.COMMITTING);
+ Instant now = Instant.now();
+ ExecutableWork stuckWork1 = createWork(createWorkItem(1L, 1L,
shardedKey1), () -> now);
+ ExecutableWork unstuckWork2 = createWork(createWorkItem(2L, 1L,
shardedKey2), () -> now);
activeWorkState.activateWorkForKey(stuckWork1);
- activeWorkState.activateWorkForKey(stuckWork2);
+ activeWorkState.activateWorkForKey(unstuckWork2);
- activeWorkState.invalidateStuckCommits(Instant.now(),
invalidatedCommits::put);
+ stuckWork1.work().setState(Work.State.COMMITTING);
+ unstuckWork2.work().setState(Work.State.PROCESSING);
- assertThat(invalidatedCommits).containsEntry(shardedKey1, stuckWork1.id());
- assertThat(invalidatedCommits).containsEntry(shardedKey2, stuckWork2.id());
- verify(computationStateCache).invalidate(shardedKey1.key(),
shardedKey1.shardingKey());
- verify(computationStateCache).invalidate(shardedKey2.key(),
shardedKey2.shardingKey());
+
assertThat(activeWorkState.hasStuckCommits(now.minus(Duration.millis(1)))).isFalse();
+ assertThat(activeWorkState.hasStuckCommits(now)).isFalse();
+
assertThat(activeWorkState.hasStuckCommits(now.plus(Duration.millis(1)))).isTrue();
}
@Test
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java
index caa25bf8309..37655152710 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java
@@ -19,12 +19,7 @@ package
org.apache.beam.runners.dataflow.worker.windmill.work.refresh;
import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.*;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import com.google.api.services.dataflow.model.MapTask;
@@ -53,9 +48,7 @@ import
org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDa
import
org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
import org.apache.beam.runners.direct.Clock;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashBasedTable;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Table;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -103,13 +96,11 @@ public class ActiveWorkRefresherTest {
private ActiveWorkRefresher createActiveWorkRefresher(
Supplier<Instant> clock,
int activeWorkRefreshPeriodMillis,
- int stuckCommitDurationMillis,
Supplier<Collection<ComputationState>> computations,
ActiveWorkRefresher.HeartbeatTracker heartbeatTracker) {
return new ActiveWorkRefresher(
clock,
activeWorkRefreshPeriodMillis,
- stuckCommitDurationMillis,
computations,
DataflowExecutionStateSampler.instance(),
Executors.newSingleThreadScheduledExecutor(),
@@ -180,7 +171,6 @@ public class ActiveWorkRefresherTest {
createActiveWorkRefresher(
fakeClock::now,
activeWorkRefreshPeriodMillis,
- 0,
() -> computations,
heartbeats -> heartbeatsSent::countDown);
@@ -238,7 +228,6 @@ public class ActiveWorkRefresherTest {
createActiveWorkRefresher(
fakeClock::now,
activeWorkRefreshPeriodMillis,
- 0,
() -> computations,
heartbeats -> heartbeatsSent::countDown);
@@ -248,64 +237,6 @@ public class ActiveWorkRefresherTest {
activeWorkRefresher.stop();
}
- @Test
- public void testInvalidateStuckCommits() throws InterruptedException {
- int stuckCommitDurationMillis = 100;
- Table<ComputationState, ExecutableWork, WindmillStateCache.ForComputation>
computations =
- HashBasedTable.create();
- WindmillStateCache stateCache =
WindmillStateCache.builder().setSizeMb(100).build();
- ByteString key = ByteString.EMPTY;
- for (int i = 0; i < 5; i++) {
- WindmillStateCache.ForComputation perComputationStateCache =
- spy(stateCache.forComputation(COMPUTATION_ID_PREFIX + i));
- ComputationState computationState = spy(createComputationState(i,
perComputationStateCache));
- ExecutableWork fakeWork = createOldWork(ShardedKey.create(key, i), i,
ignored -> {});
- fakeWork.work().setState(Work.State.COMMITTING);
- computationState.activateWork(fakeWork);
- computations.put(computationState, fakeWork, perComputationStateCache);
- }
-
- TestClock fakeClock = new TestClock(Instant.now());
- CountDownLatch invalidateStuckCommitRan = new
CountDownLatch(computations.size());
-
- // Count down the latch every time to avoid waiting/sleeping arbitrarily.
- for (ComputationState computation : computations.rowKeySet()) {
- doAnswer(
- invocation -> {
- invocation.callRealMethod();
- invalidateStuckCommitRan.countDown();
- return null;
- })
- .when(computation)
- .invalidateStuckCommits(any(Instant.class));
- }
-
- ActiveWorkRefresher activeWorkRefresher =
- createActiveWorkRefresher(
- fakeClock::now,
- 0,
- stuckCommitDurationMillis,
- computations.rowMap()::keySet,
- ignored -> () -> {});
-
- activeWorkRefresher.start();
- fakeClock.advance(Duration.millis(stuckCommitDurationMillis));
- invalidateStuckCommitRan.await();
- activeWorkRefresher.stop();
-
- for (Table.Cell<ComputationState, ExecutableWork,
WindmillStateCache.ForComputation> cell :
- computations.cellSet()) {
- ComputationState computation = cell.getRowKey();
- ExecutableWork work = cell.getColumnKey();
- WindmillStateCache.ForComputation perComputationStateCache =
cell.getValue();
- verify(perComputationStateCache, times(1))
- .invalidate(eq(key), eq(work.getWorkItem().getShardingKey()));
- verify(computation, times(1))
- .completeWorkAndScheduleNextWorkForKey(
- eq(ShardedKey.create(key, work.getWorkItem().getShardingKey())),
eq(work.id()));
- }
- }
-
static class TestClock implements Clock {
private Instant time;
