Skip to content

Commit d032f0e

Browse files
lalitiumlalitium
authored
Add drain support for Dataflow and Flink (#38786)
* Add drain support for Dataflow and Flink * Add PipelineResult drain API
1 parent 0bc7852 commit d032f0e

6 files changed

Lines changed: 255 additions & 29 deletions

File tree

‎runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java‎

Lines changed: 56 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@
1818
package org.apache.beam.runners.flink;
1919

2020
import java.io.IOException;
21+
import java.util.concurrent.CompletableFuture;
2122
import java.util.concurrent.ExecutionException;
2223
import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
2324
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
2425
import org.apache.beam.sdk.PipelineResult;
2526
import org.apache.beam.sdk.metrics.MetricResults;
2627
import org.apache.flink.api.common.JobStatus;
2728
import org.apache.flink.core.execution.JobClient;
29+
import org.apache.flink.core.execution.SavepointFormatType;
30+
import org.checkerframework.checker.nullness.qual.Nullable;
2831
import org.joda.time.Duration;
2932
import org.slf4j.Logger;
3033
import org.slf4j.LoggerFactory;
@@ -39,6 +42,7 @@ public class FlinkDetachedRunnerResult implements PipelineResult {
3942

4043
private JobClient jobClient;
4144
private int jobCheckIntervalInSecs;
45+
private volatile @Nullable CompletableFuture<String> drainSavepointFuture;
4246

4347
FlinkDetachedRunnerResult(JobClient jobClient, int jobCheckIntervalInSecs) {
4448
this.jobClient = jobClient;
@@ -47,10 +51,25 @@ public class FlinkDetachedRunnerResult implements PipelineResult {
4751

4852
@Override
4953
public State getState() {
54+
CompletableFuture<String> drainFuture = drainSavepointFuture;
55+
if (drainFuture != null) {
56+
try {
57+
return getDrainState(drainFuture);
58+
} catch (IOException e) {
59+
LOG.warn("Failed to drain Flink job. Querying Flink job state instead.", e);
60+
}
61+
}
62+
return getFlinkJobState();
63+
}
64+
65+
private State getFlinkJobState() {
5066
try {
5167
return toBeamJobState(jobClient.getJobStatus().get());
52-
} catch (InterruptedException | ExecutionException e) {
53-
throw new RuntimeException("Fail to get flink job state", e);
68+
} catch (InterruptedException e) {
69+
Thread.currentThread().interrupt();
70+
throw new RuntimeException("Failed to get Flink job state", e);
71+
} catch (ExecutionException e) {
72+
throw new RuntimeException("Failed to get Flink job state", e);
5473
}
5574
}
5675

@@ -66,7 +85,11 @@ private MetricsContainerStepMap getMetricsContainerStepMap() {
6685
.getAccumulators()
6786
.get()
6887
.getOrDefault(FlinkMetricContainer.ACCUMULATOR_NAME, new MetricsContainerStepMap());
69-
} catch (InterruptedException | ExecutionException e) {
88+
} catch (InterruptedException e) {
89+
Thread.currentThread().interrupt();
90+
LOG.warn("Fail to get flink job accumulators", e);
91+
return new MetricsContainerStepMap();
92+
} catch (ExecutionException e) {
7093
LOG.warn("Fail to get flink job accumulators", e);
7194
return new MetricsContainerStepMap();
7295
}
@@ -76,12 +99,40 @@ private MetricsContainerStepMap getMetricsContainerStepMap() {
7699
public State cancel() throws IOException {
77100
try {
78101
this.jobClient.cancel().get();
79-
} catch (InterruptedException | ExecutionException e) {
102+
} catch (InterruptedException e) {
103+
Thread.currentThread().interrupt();
104+
throw new RuntimeException("Fail to cancel flink job", e);
105+
} catch (ExecutionException e) {
80106
throw new RuntimeException("Fail to cancel flink job", e);
81107
}
82108
return getState();
83109
}
84110

111+
@Override
112+
public synchronized State drain() throws IOException {
113+
CompletableFuture<String> drainFuture = drainSavepointFuture;
114+
if (drainFuture == null || drainFuture.isCompletedExceptionally()) {
115+
drainFuture = this.jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT);
116+
drainSavepointFuture = drainFuture;
117+
}
118+
return getDrainState(drainFuture);
119+
}
120+
121+
private State getDrainState(CompletableFuture<String> drainFuture) throws IOException {
122+
if (!drainFuture.isDone()) {
123+
return State.RUNNING;
124+
}
125+
try {
126+
drainFuture.get();
127+
return State.DONE;
128+
} catch (InterruptedException e) {
129+
Thread.currentThread().interrupt();
130+
throw new IOException("Failed to drain Flink job", e);
131+
} catch (ExecutionException e) {
132+
throw new IOException("Failed to drain Flink job", e.getCause());
133+
}
134+
}
135+
85136
@Override
86137
public State waitUntilFinish() {
87138
return waitUntilFinish(Duration.millis(Long.MAX_VALUE));
@@ -100,6 +151,7 @@ public State waitUntilFinish(Duration duration) {
100151
try {
101152
Thread.sleep(jobCheckIntervalInSecs * 1000L);
102153
} catch (InterruptedException e) {
154+
Thread.currentThread().interrupt();
103155
throw new RuntimeException(e);
104156
}
105157
}

‎runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunnerResult.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,12 @@ public State cancel() {
6464
return State.DONE;
6565
}
6666

67+
@Override
68+
public State drain() {
69+
// We can only be called here when we are done.
70+
return State.DONE;
71+
}
72+
6773
@Override
6874
public State waitUntilFinish() {
6975
return State.DONE;

‎runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,20 @@
1919

2020
import static org.hamcrest.MatcherAssert.assertThat;
2121
import static org.hamcrest.core.Is.is;
22+
import static org.junit.Assert.assertSame;
23+
import static org.junit.Assert.fail;
24+
import static org.mockito.Mockito.mock;
25+
import static org.mockito.Mockito.times;
26+
import static org.mockito.Mockito.verify;
27+
import static org.mockito.Mockito.when;
2228

29+
import java.io.IOException;
2330
import java.util.Collections;
31+
import java.util.concurrent.CompletableFuture;
2432
import org.apache.beam.sdk.PipelineResult;
33+
import org.apache.flink.api.common.JobStatus;
34+
import org.apache.flink.core.execution.JobClient;
35+
import org.apache.flink.core.execution.SavepointFormatType;
2536
import org.joda.time.Duration;
2637
import org.junit.Test;
2738

@@ -47,4 +58,85 @@ public void testCancelDoesNotThrowAnException() {
4758
result.cancel();
4859
assertThat(result.getState(), is(PipelineResult.State.DONE));
4960
}
61+
62+
@Test
63+
public void testDrainDoneResultDoesNotThrowAnException() throws Exception {
64+
FlinkRunnerResult result = new FlinkRunnerResult(Collections.emptyMap(), 100);
65+
assertThat(result.drain(), is(PipelineResult.State.DONE));
66+
}
67+
68+
@Test
69+
public void testDetachedDrainReturnsRunningThenDone() throws Exception {
70+
JobClient jobClient = mock(JobClient.class);
71+
CompletableFuture<String> drainFuture = new CompletableFuture<>();
72+
when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT))
73+
.thenReturn(drainFuture);
74+
FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1);
75+
76+
assertThat(result.drain(), is(PipelineResult.State.RUNNING));
77+
assertThat(result.getState(), is(PipelineResult.State.RUNNING));
78+
79+
drainFuture.complete("savepoint");
80+
assertThat(result.getState(), is(PipelineResult.State.DONE));
81+
verify(jobClient).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT);
82+
}
83+
84+
@Test
85+
public void testDetachedDrainFailureThrowsIOException() throws Exception {
86+
JobClient jobClient = mock(JobClient.class);
87+
CompletableFuture<String> drainFuture = new CompletableFuture<>();
88+
RuntimeException failure = new RuntimeException("savepoint failed");
89+
drainFuture.completeExceptionally(failure);
90+
when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT))
91+
.thenReturn(drainFuture);
92+
FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1);
93+
94+
try {
95+
result.drain();
96+
fail("Expected IOException");
97+
} catch (IOException e) {
98+
assertThat(e.getMessage(), is("Failed to drain Flink job"));
99+
assertSame(failure, e.getCause());
100+
}
101+
}
102+
103+
@Test
104+
public void testDetachedGetStateFallsBackAfterDrainFailure() throws Exception {
105+
JobClient jobClient = mock(JobClient.class);
106+
CompletableFuture<String> drainFuture = new CompletableFuture<>();
107+
drainFuture.completeExceptionally(new RuntimeException("savepoint failed"));
108+
when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT))
109+
.thenReturn(drainFuture);
110+
when(jobClient.getJobStatus()).thenReturn(CompletableFuture.completedFuture(JobStatus.RUNNING));
111+
FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1);
112+
113+
try {
114+
result.drain();
115+
fail("Expected IOException");
116+
} catch (IOException expected) {
117+
assertThat(result.getState(), is(PipelineResult.State.RUNNING));
118+
}
119+
}
120+
121+
@Test
122+
public void testDetachedDrainRetriesAfterFailure() throws Exception {
123+
JobClient jobClient = mock(JobClient.class);
124+
CompletableFuture<String> failedDrainFuture = new CompletableFuture<>();
125+
failedDrainFuture.completeExceptionally(new RuntimeException("savepoint failed"));
126+
CompletableFuture<String> retryDrainFuture = new CompletableFuture<>();
127+
when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT))
128+
.thenReturn(failedDrainFuture, retryDrainFuture);
129+
FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1);
130+
131+
try {
132+
result.drain();
133+
fail("Expected IOException");
134+
} catch (IOException expected) {
135+
assertThat(result.drain(), is(PipelineResult.State.RUNNING));
136+
}
137+
138+
retryDrainFuture.complete("savepoint");
139+
assertThat(result.getState(), is(PipelineResult.State.DONE));
140+
verify(jobClient, times(2)).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT);
141+
}
50142
}

‎runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java‎

Lines changed: 42 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -432,70 +432,87 @@ private Exception processJobMessages(
432432
}
433433

434434
private AtomicReference<FutureTask<State>> cancelState = new AtomicReference<>();
435+
private AtomicReference<FutureTask<State>> drainState = new AtomicReference<>();
435436

436437
@SuppressWarnings("Slf4jFormatShouldBeConst")
437438
@Override
438439
public State cancel() throws IOException {
439-
// Enforce that a cancel() call on the job is done at most once - as
440-
// a workaround for Dataflow service's current bugs with multiple
441-
// cancellation, where it may sometimes return an error when cancelling
442-
// a job that was already cancelled, but still report the job state as
443-
// RUNNING.
444-
// To partially work around these issues, we absorb duplicate cancel()
445-
// calls. This, of course, doesn't address the case when the job terminates
446-
// externally almost concurrently to calling cancel(), but at least it
447-
// makes it possible to safely call cancel() multiple times and from
448-
// multiple threads in one program.
449-
FutureTask<State> tentativeCancelTask =
440+
return requestJobState(cancelState, "JOB_STATE_CANCELLED", "cancel", "Cancel");
441+
}
442+
443+
@Override
444+
public State drain() throws IOException {
445+
return requestJobState(drainState, "JOB_STATE_DRAINED", "drain", "Drain");
446+
}
447+
448+
@SuppressWarnings("Slf4jFormatShouldBeConst")
449+
private State requestJobState(
450+
AtomicReference<FutureTask<State>> requestedState,
451+
String dataflowRequestedState,
452+
String action,
453+
String capitalizedAction)
454+
throws IOException {
455+
// Enforce that a lifecycle request on the job is done at most once. This preserves the
456+
// existing cancel() behavior and keeps duplicate drain() calls idempotent from one client.
457+
FutureTask<State> tentativeTask =
450458
new FutureTask<>(
451459
() -> {
452460
Job content = new Job();
453461
content.setProjectId(getProjectId());
454462
String currentJobId = getJobId();
455463
content.setId(currentJobId);
456-
content.setRequestedState("JOB_STATE_CANCELLED");
464+
content.setRequestedState(dataflowRequestedState);
457465
try {
458466
Job job = dataflowClient.updateJob(currentJobId, content);
459467
return MonitoringUtil.toState(job.getCurrentState());
460468
} catch (IOException e) {
461469
State state = getState();
470+
String message = e.getMessage();
462471
if (state.isTerminal()) {
463-
LOG.warn("Cancel failed because job is already terminated. State is {}", state);
472+
LOG.warn(
473+
"{} failed because job is already terminated. State is {}",
474+
capitalizedAction,
475+
state);
464476
return state;
465-
} else if (e.getMessage().contains("has terminated")) {
477+
} else if (message != null && message.contains("has terminated")) {
466478
// This handles the case where the getState() call above returns RUNNING but the
467-
// cancel was rejected because the job is in fact done. Hopefully, someday we can
479+
// request was rejected because the job is in fact done. Hopefully, someday we can
468480
// delete this code if there is better consistency between the State and whether
469-
// Cancel succeeds.
481+
// lifecycle requests succeed.
470482
//
471483
// Example message:
472484
// Workflow modification failed. Causes: (7603adc9e9bff51e): Cannot perform
473485
// operation 'cancel' on Job: 2017-04-01_22_50_59-9269855660514862348. Job has
474486
// terminated in state SUCCESS: Workflow job:
475487
// 2017-04-01_22_50_59-9269855660514862348 succeeded.
476-
LOG.warn("Cancel failed because job is already terminated.", e);
488+
LOG.warn("{} failed because job is already terminated.", capitalizedAction, e);
477489
return state;
478490
} else {
479491
String errorMsg =
480492
String.format(
481-
"Failed to cancel job in state %s, "
482-
+ "please go to the Developers Console to cancel it manually: %s",
493+
"Failed to %s job in state %s, "
494+
+ "please go to the Developers Console to %s it manually: %s",
495+
action,
483496
state,
497+
action,
484498
MonitoringUtil.getJobMonitoringPageURL(
485499
getProjectId(), getRegion(), getJobId()));
486500
LOG.warn(errorMsg);
487501
throw new IOException(errorMsg, e);
488502
}
489503
}
490504
});
491-
if (cancelState.compareAndSet(null, tentativeCancelTask)) {
492-
// This thread should perform cancellation, while others will
493-
// only wait for the result.
494-
cancelState.get().run();
505+
if (requestedState.compareAndSet(null, tentativeTask)) {
506+
// This thread should perform the lifecycle request, while others will only wait for the
507+
// result.
508+
requestedState.get().run();
495509
}
496510
try {
497-
return cancelState.get().get();
498-
} catch (InterruptedException | ExecutionException e) {
511+
return requestedState.get().get();
512+
} catch (InterruptedException e) {
513+
Thread.currentThread().interrupt();
514+
throw new IOException(e);
515+
} catch (ExecutionException e) {
499516
throw new IOException(e);
500517
}
501518
}

0 commit comments

Comments
 (0)