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

Example 1 with WorkItemServiceState

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

the class DataflowWorkProgressUpdater method reportProgressHelper.

@Override
protected void reportProgressHelper() throws Exception {
    if (wasAskedToAbort) {
        LOG.info("Service already asked to abort work item, not reporting ignored progress.");
        return;
    }
    WorkItemServiceState result = workItemStatusClient.reportUpdate(dynamicSplitResultToReport, Duration.millis(requestedLeaseDurationMs));
    if (result != null) {
        if (result.getCompleteWorkStatus() != null && result.getCompleteWorkStatus().getCode() != com.google.rpc.Code.OK.getNumber()) {
            LOG.info("Service asked worker to abort with status: {}", result.getCompleteWorkStatus());
            wasAskedToAbort = true;
            worker.abort();
            return;
        }
        if (result.getHotKeyDetection() != null && result.getHotKeyDetection().getUserStepName() != null) {
            HotKeyDetection hotKeyDetection = result.getHotKeyDetection();
            // which is the correct key. The key is also translated into a Java object in the reader.
            if (options.isHotKeyLoggingEnabled()) {
                hotKeyLogger.logHotKeyDetection(hotKeyDetection.getUserStepName(), TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()), workItemStatusClient.getExecutionContext().getKey());
            } else {
                hotKeyLogger.logHotKeyDetection(hotKeyDetection.getUserStepName(), TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()));
            }
        }
        // Resets state after a successful progress report.
        dynamicSplitResultToReport = null;
        progressReportIntervalMs = nextProgressReportInterval(fromCloudDuration(result.getReportStatusInterval()).getMillis(), leaseRemainingTime(getLeaseExpirationTimestamp(result)));
        ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest();
        if (suggestedStopPoint != null) {
            LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint);
            dynamicSplitResultToReport = worker.requestDynamicSplit(SourceTranslationUtils.toDynamicSplitRequest(suggestedStopPoint));
        }
    }
}

Example 2 with WorkItemServiceState

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

the class CounterShortIdCache method storeNewShortIds.

/**
 * Add any new short ids received to the table. The outgoing request will have the full counter
 * updates, and the incoming responses have the associated short ids. By matching up short ids
 * with the counters in order we can build a mapping of name -> short_id for future use.
 */
public void storeNewShortIds(final ReportWorkItemStatusRequest request, final ReportWorkItemStatusResponse reply) {
    checkArgument(request.getWorkItemStatuses() != null && reply.getWorkItemServiceStates() != null && request.getWorkItemStatuses().size() == reply.getWorkItemServiceStates().size(), "RequestWorkItemStatus request and response are unbalanced, status: %s, states: %s", request.getWorkItemStatuses(), reply.getWorkItemServiceStates());
    for (int i = 0; i < request.getWorkItemStatuses().size(); i++) {
        WorkItemServiceState state = reply.getWorkItemServiceStates().get(i);
        WorkItemStatus status = request.getWorkItemStatuses().get(i);
        if (state.getMetricShortId() == null) {
            continue;
        }
        checkArgument(status.getCounterUpdates() != null, "Response has shortids but no corresponding CounterUpdate");
        for (MetricShortId shortIdMsg : state.getMetricShortId()) {
            int metricIndex = MoreObjects.firstNonNull(shortIdMsg.getMetricIndex(), 0);
            checkArgument(metricIndex < status.getCounterUpdates().size(), "Received aggregate index outside range of sent update %s >= %s", shortIdMsg.getMetricIndex(), status.getCounterUpdates().size());
            CounterUpdate update = status.getCounterUpdates().get(metricIndex);
            cache.insert(update, checkNotNull(shortIdMsg.getShortId(), "Shortid should be non-null"));
        }
    }
}

Example 3 with WorkItemServiceState

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

the class DataflowWorkProgressUpdaterTest method workProgressSendsSplitResults.

/**
 * Verifies that the update after a split is requested acknowledeges it.
 */
@Test
public void workProgressSendsSplitResults() 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));
    // The second update should be sent at startTime + 1300 ms and includes the split response.
    // Also schedule another update after 1000ms.
    when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class))).thenReturn(generateServiceState(null, 1000));
    executor.runNextRunnable();
    assertEquals(clock.currentTimeMillis(), startTime + 1300);
    // Verify that the update includes the respuonse to the split request.
    verify(workItemStatusClient, atLeastOnce()).reportUpdate(splitResultCaptor.capture(), isA(Duration.class));
    assertEquals("Second update is sent and contains the latest split result", splitResultCaptor.getValue(), new NativeReader.DynamicSplitResultWithPosition(cloudPositionToReaderPosition(ReaderTestUtils.positionAtIndex(2L))));
    executor.runNextRunnable();
    verify(workItemStatusClient, times(3)).reportUpdate(splitResultCaptor.capture(), isA(Duration.class));
    // Stop the progressUpdater now, and expect the last update immediately
    progressUpdater.stopReportingProgress();
}

Example 4 with WorkItemServiceState

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

the class WorkItemStatusClient method execute.

@Nullable
private synchronized WorkItemServiceState execute(WorkItemStatus status) throws IOException {
    WorkItemServiceState result = workUnitClient.reportWorkItemStatus(status);
    if (result != null) {
        nextReportIndex = result.getNextReportIndex();
        if (nextReportIndex == null && !status.getCompleted()) {
            LOG.error("Missing next work index in {} when reporting {}.", result, status);
        }
        commitMetrics();
    }
    if (status.getCompleted()) {
        checkState(!finalStateSent, "cannot reportUpdates after sending a final state");
        finalStateSent = true;
    }
    return result;
}

Example 5 with WorkItemServiceState

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

the class DataflowWorkUnitClient method reportWorkItemStatus.

/**
 * Reports the status of the most recently requested work item.
 */
@Override
public WorkItemServiceState reportWorkItemStatus(WorkItemStatus workItemStatus) throws IOException {
    DateTime endTime = DateTime.now();
    workItemStatus.setFactory(Transport.getJsonFactory());
    logger.debug("Reporting work status: {}", workItemStatus);
    // in the event this status is associated with a dummy work item.
    if (firstNonNull(workItemStatus.getCompleted(), Boolean.FALSE) && DataflowWorkerLoggingMDC.getStageName() != null) {
        DateTime startTime = stageStartTime.get();
        if (startTime != null) {
            // elapsed time can be negative by time correction
            long elapsed = endTime.getMillis() - startTime.getMillis();
            int numErrors = workItemStatus.getErrors() == null ? 0 : workItemStatus.getErrors().size();
            // This thread should have been tagged with the stage start time during getWorkItem(),
            logger.info("Finished processing stage {} with {} errors in {} seconds ", DataflowWorkerLoggingMDC.getStageName(), numErrors, (double) elapsed / 1000);
        }
    }
    shortIdCache.shortenIdsIfAvailable(workItemStatus.getCounterUpdates());
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest().setWorkerId(options.getWorkerId()).setWorkItemStatuses(Collections.singletonList(workItemStatus)).setCurrentWorkerTime(toCloudTime(endTime));
    ReportWorkItemStatusResponse result = dataflow.projects().locations().jobs().workItems().reportStatus(options.getProject(), options.getRegion(), options.getJobId(), request).execute();
    if (result == null) {
        logger.warn("Report work item status response: null");
        throw new IOException("Got null work item status response");
    }
    if (result.getWorkItemServiceStates() == null) {
        logger.warn("Report work item status response: {}", result);
        throw new IOException("Report work item status contained no work item service states");
    }
    if (result.getWorkItemServiceStates().size() != 1) {
        logger.warn("Report work item status response: {}", result);
        throw new IOException("This version of the SDK expects exactly one work item service state from the service " + "but got " + result.getWorkItemServiceStates().size() + " states");
    }
    shortIdCache.storeNewShortIds(request, result);
    WorkItemServiceState state = result.getWorkItemServiceStates().get(0);
    logger.debug("ReportWorkItemStatus result: {}", state);
    return state;
}