DataflowPipelineJob implements PipelineResult { @Override public State cancel() throws IOException { FutureTask<State> tentativeCancelTask = new FutureTask<>( () -> { Job content = new Job(); content.setProjectId(getProjectId()); String currentJobId = getJobId(); content.setId(currentJobId); content.setRequestedState("JOB_STATE_CANCELLED"); try { Job job = dataflowClient.updateJob(currentJobId, content); return MonitoringUtil.toState(job.getCurrentState()); } catch (IOException e) { State state = getState(); if (state.isTerminal()) { LOG.warn("Cancel failed because job is already terminated. State is {}", state); return state; } else if (e.getMessage().contains("has terminated")) { LOG.warn("Cancel failed because job is already terminated.", e); return state; } else { String errorMsg = String.format( "Failed to cancel job in state %s, " + "please go to the Developers Console to cancel it manually: %s", state, MonitoringUtil.getJobMonitoringPageURL( getProjectId(), getRegion(), getJobId())); LOG.warn(errorMsg); throw new IOException(errorMsg, e); } } }); if (cancelState.compareAndSet(null, tentativeCancelTask)) { cancelState.get().run(); } try { return cancelState.get().get(); } catch (InterruptedException | ExecutionException e) { throw new IOException(e); } } DataflowPipelineJob(
      DataflowClient dataflowClient,
      String jobId,
      DataflowPipelineOptions dataflowOptions,
      Map<AppliedPTransform<?, ?, ?>, String> transformStepNames); String getJobId(); String getProjectId(); DataflowPipelineOptions getDataflowOptions(); String getRegion(); DataflowPipelineJob getReplacedByJob(); @Override @Nullable State waitUntilFinish(); @Override @Nullable State waitUntilFinish(Duration duration); @Nullable @VisibleForTesting State waitUntilFinish(Duration duration, MonitoringUtil.JobMessagesHandler messageHandler); @Override State cancel(); @Override State getState(); @Override MetricResults metrics();  }
@Test public void testCancelUnterminatedJobThatSucceeds() throws IOException { Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenReturn(new Job().setCurrentState("JOB_STATE_CANCELLED")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); assertEquals(State.CANCELLED, job.cancel()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); content.setRequestedState("JOB_STATE_CANCELLED"); verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content)); verifyNoMoreInteractions(mockJobs); }
@Test public void testCancelUnterminatedJobThatFails() throws IOException { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_RUNNING"); when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException("Some random IOException")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); thrown.expect(IOException.class); thrown.expectMessage( "Failed to cancel job in state RUNNING, " + "please go to the Developers Console to cancel it manually:"); job.cancel(); }
@Test public void testCancelTerminatedJobWithStaleState() throws IOException { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_RUNNING"); when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException("Job has terminated in state SUCCESS")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); State returned = job.cancel(); assertThat(returned, equalTo(State.RUNNING)); expectedLogs.verifyWarn("Cancel failed because job is already terminated."); }
@Test public void testCancelTerminatedJob() throws IOException { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_FAILED"); when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException()); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); assertEquals(State.FAILED, job.cancel()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); content.setRequestedState("JOB_STATE_CANCELLED"); verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content)); verify(mockJobs).get(PROJECT_ID, REGION_ID, JOB_ID); verifyNoMoreInteractions(mockJobs); }