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

Example 1 with WorkItemStatus

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

the class BatchDataflowWorkerTest method testWhenNoWorkIsReturnedThatWeImmediatelyRetry.

@Test
public void testWhenNoWorkIsReturnedThatWeImmediatelyRetry() throws Exception {
    final String workItemId = "14";
    BatchDataflowWorker worker = new BatchDataflowWorker(null, /* pipeline */
    SdkHarnessRegistries.emptySdkHarnessRegistry(), mockWorkUnitClient, IntrinsicMapTaskExecutorFactory.defaultFactory(), options);
    WorkItem workItem = new WorkItem();
    workItem.setId(Long.parseLong(workItemId));
    workItem.setJobId("SuccessfulEmptyMapTask");
    workItem.setInitialReportIndex(12L);
    workItem.setMapTask(new MapTask().setInstructions(new ArrayList<ParallelInstruction>()).setStageName("testStage"));
    workItem.setLeaseExpireTime(TimeUtil.toCloudTime(Instant.now()));
    workItem.setReportStatusInterval(TimeUtil.toCloudDuration(Duration.standardMinutes(1)));
    when(mockWorkUnitClient.getWorkItem()).thenReturn(Optional.<WorkItem>absent()).thenReturn(Optional.of(workItem));
    assertTrue(worker.getAndPerformWork());
    verify(mockWorkUnitClient).reportWorkItemStatus(MockitoHamcrest.argThat(new TypeSafeMatcher<WorkItemStatus>() {

        @Override
        public void describeTo(Description description) {
        }

        @Override
        protected boolean matchesSafely(WorkItemStatus item) {
            assertTrue(item.getCompleted());
            assertEquals(workItemId, item.getWorkItemId());
            return true;
        }
    }));
}

Example 2 with WorkItemStatus

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

the class CounterShortIdCacheTest method createWorkStatusStructuredName.

private List<WorkItemStatus> createWorkStatusStructuredName(String[]... counterNames) {
    List<WorkItemStatus> statuses = new ArrayList<>();
    for (String[] names : counterNames) {
        WorkItemStatus status = new WorkItemStatus();
        List<CounterUpdate> counterList = new ArrayList<>();
        for (String name : names) {
            counterList.add(createMetricUpdateStructuredName(name));
        }
        status.setCounterUpdates(counterList);
        statuses.add(status);
    }
    return statuses;
}

Example 3 with WorkItemStatus

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

the class CounterShortIdCacheTest method createWorkStatusNameAndKind.

private List<WorkItemStatus> createWorkStatusNameAndKind(String[]... counterNames) {
    List<WorkItemStatus> statuses = new ArrayList<>();
    for (String[] names : counterNames) {
        WorkItemStatus status = new WorkItemStatus();
        List<CounterUpdate> counterList = new ArrayList<>();
        for (String name : names) {
            counterList.add(createMetricUpdateNameAndKind(name));
        }
        status.setCounterUpdates(counterList);
        statuses.add(status);
    }
    return statuses;
}

Example 4 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus 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 5 with WorkItemStatus

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