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

Example 1 with ReportWorkItemStatusRequest

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

the class CounterShortIdCacheTest method testValidateNumberStatusesAndStates.

@Test
public void testValidateNumberStatusesAndStates() {
    CounterShortIdCache cache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    request.setWorkItemStatuses(createWorkStatusNameAndKind(new String[] { "counter" }, new String[] { "counter2" }));
    reply.setWorkItemServiceStates(createWorkServiceState(new MetricShortId[] { createMetricShortId(0, 1000L) }));
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("RequestWorkItemStatus request and response are unbalanced");
    cache.storeNewShortIds(request, reply);
}

Example 2 with ReportWorkItemStatusRequest

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

the class CounterShortIdCacheTest method testValidateAggregateIndexOutOfRange.

@Test
public void testValidateAggregateIndexOutOfRange() {
    CounterShortIdCache cache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    request.setWorkItemStatuses(createWorkStatusNameAndKind(new String[] { "counter" }));
    reply.setWorkItemServiceStates(createWorkServiceState(new MetricShortId[] { createMetricShortId(1000, 1000L) }));
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Received aggregate index outside range of sent update");
    cache.storeNewShortIds(request, reply);
}

Example 3 with ReportWorkItemStatusRequest

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

the class CounterShortIdCacheTest method testCacheNameAndKind.

@Test
public void testCacheNameAndKind() {
    CounterShortIdCache shortIdCache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    // setup mock counters, three work statuses, one with two counters, one with one, one with none
    request.setWorkItemStatuses(createWorkStatusNameAndKind(new String[] { "counter", "counter1" }, new String[] {}, new String[] { "counter2" }));
    reply.setWorkItemServiceStates(createWorkServiceState(new Long[] { 1000L, 1001L }, new Long[] {}, new Long[] { 1002L }));
    // Verify the empty case
    WorkItemStatus status1 = request.getWorkItemStatuses().get(0);
    WorkItemStatus status2 = request.getWorkItemStatuses().get(1);
    WorkItemStatus status3 = request.getWorkItemStatuses().get(2);
    shortIdCache.shortenIdsIfAvailable(status1.getCounterUpdates());
    for (CounterUpdate update : status1.getCounterUpdates()) {
        assertNull(update.getShortId());
    }
    // Add the shortIds
    shortIdCache.storeNewShortIds(request, reply);
    shortIdCache.shortenIdsIfAvailable(status1.getCounterUpdates());
    shortIdCache.shortenIdsIfAvailable(status2.getCounterUpdates());
    shortIdCache.shortenIdsIfAvailable(status3.getCounterUpdates());
    checkStatusAndShortIds(status1, 1000L, 1001L);
    checkStatusAndShortIds(status2);
    checkStatusAndShortIds(status3, 1002L);
}

Example 4 with ReportWorkItemStatusRequest

use of com.google.api.services.dataflow.model.ReportWorkItemStatusRequest 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;
}

Example 5 with ReportWorkItemStatusRequest

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

the class CounterShortIdCacheTest method testValidateShortIdsButNoUpdate.

@Test
public void testValidateShortIdsButNoUpdate() {
    CounterShortIdCache cache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    request.setWorkItemStatuses(Arrays.asList(new WorkItemStatus()));
    reply.setWorkItemServiceStates(createWorkServiceState(new Long[] { 1000L }));
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Response has shortids but no corresponding CounterUpdate");
    cache.storeNewShortIds(request, reply);
}