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