Examples with WorkItemServiceState com.google.api.services.dataflow.model.WorkItemServiceState used on opensource projects

Example 6 with WorkItemServiceState

use of com.google.api.services.dataflow.model.WorkItemServiceState in project beam by apache.

the class DataflowWorkProgressUpdaterTest method generateServiceAbort.

private WorkItemServiceState generateServiceAbort() {
    WorkItemServiceState responseState = new WorkItemServiceState();
    responseState.setCompleteWorkStatus(new Status().setCode(com.google.rpc.Code.ABORTED_VALUE).setMessage("Worker was asked to abort!"));
    return responseState;
}

Example 7 with WorkItemServiceState

use of com.google.api.services.dataflow.model.WorkItemServiceState in project beam by apache.

the class DataflowWorkProgressUpdaterTest method generateServiceState.

private WorkItemServiceState generateServiceState(@Nullable Position suggestedStopPosition, long millisToNextUpdate) {
    WorkItemServiceState responseState = new WorkItemServiceState();
    responseState.setFactory(Transport.getJsonFactory());
    responseState.setLeaseExpireTime(toCloudTime(new Instant(clock.currentTimeMillis() + LEASE_MS)));
    responseState.setReportStatusInterval(toCloudDuration(Duration.millis(millisToNextUpdate)));
    if (suggestedStopPosition != null) {
        responseState.setSplitRequest(ReaderTestUtils.approximateSplitRequestAtPosition(suggestedStopPosition));
    }
    HotKeyDetection hotKeyDetection = new HotKeyDetection();
    hotKeyDetection.setUserStepName(STEP_ID);
    hotKeyDetection.setHotKeyAge(toCloudDuration(HOT_KEY_AGE));
    responseState.setHotKeyDetection(hotKeyDetection);
    return responseState;
}

Example 8 with WorkItemServiceState

use of com.google.api.services.dataflow.model.WorkItemServiceState in project beam by apache.

the class DataflowWorkProgressUpdaterTest method workProgressUpdaterSendsLastPendingUpdateWhenStopped.

/**
 * Verifies that a last update is sent when there is an unacknowledged split request.
 */
@Test
public void workProgressUpdaterSendsLastPendingUpdateWhenStopped() throws Exception {
    // The setup process sends one update after 300ms. Enqueue another that should be scheduled
    // 1000ms after that.
    WorkItemServiceState firstResponse = generateServiceState(ReaderTestUtils.positionAtIndex(2L), 1000);
    when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class))).thenReturn(firstResponse);
    when(worker.getWorkerProgress()).thenReturn(cloudProgressToReaderProgress(ReaderTestUtils.approximateProgressAtIndex(1L)));
    when(worker.requestDynamicSplit(toDynamicSplitRequest(firstResponse.getSplitRequest()))).thenReturn(new NativeReader.DynamicSplitResultWithPosition(cloudPositionToReaderPosition(firstResponse.getSplitRequest().getPosition())));
    progressUpdater.startReportingProgress();
    executor.runNextRunnable();
    // The initial update should be sent at startTime + 300 ms.
    assertEquals(clock.currentTimeMillis(), startTime + 300);
    verify(workItemStatusClient).reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class));
    verify(worker).requestDynamicSplit(isA(DynamicSplitRequest.class));
    clock.setTime(clock.currentTimeMillis() + 100);
    // Stop the progressUpdater now, and expect the last update immediately
    progressUpdater.stopReportingProgress();
    assertEquals(clock.currentTimeMillis(), startTime + 400);
    // Verify that the last update is sent immediately and contained the latest split result.
    verify(workItemStatusClient, atLeastOnce()).reportUpdate(splitResultCaptor.capture(), isA(Duration.class));
    assertEquals("Final update is sent and contains the latest split result", splitResultCaptor.getValue(), new NativeReader.DynamicSplitResultWithPosition(cloudPositionToReaderPosition(ReaderTestUtils.positionAtIndex(2L))));
    // And nothing happened after that.
    verify(workItemStatusClient, Mockito.atLeastOnce()).uniqueWorkId();
    verify(workItemStatusClient, Mockito.atLeastOnce()).getExecutionContext();
    verifyNoMoreInteractions(workItemStatusClient);
}

Example 9 with WorkItemServiceState

use of com.google.api.services.dataflow.model.WorkItemServiceState in project beam by apache.

the class CounterShortIdCacheTest method createWorkServiceState.

private List<WorkItemServiceState> createWorkServiceState(Long[]... counterIds) {
    List<WorkItemServiceState> states = new ArrayList<>();
    for (Long[] ids : counterIds) {
        WorkItemServiceState state = new WorkItemServiceState();
        List<MetricShortId> shortIds = new ArrayList<>();
        for (int i = 0; i < ids.length; i++) {
            shortIds.add(createMetricShortId(i, ids[i]));
        }
        state.setMetricShortId(shortIds);
        states.add(state);
    }
    return states;
}

Example 10 with WorkItemServiceState

use of com.google.api.services.dataflow.model.WorkItemServiceState in project beam by apache.

the class CounterShortIdCacheTest method createWorkServiceState.

private List<WorkItemServiceState> createWorkServiceState(MetricShortId[]... counterIds) {
    List<WorkItemServiceState> states = new ArrayList<>();
    for (MetricShortId[] ids : counterIds) {
        WorkItemServiceState state = new WorkItemServiceState();
        List<MetricShortId> shortIds = new ArrayList<>();
        for (int i = 0; i < ids.length; i++) {
            shortIds.add(ids[i]);
        }
        state.setMetricShortId(shortIds);
        states.add(state);
    }
    return states;
}