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

Example 1 with ApproximateSplitRequest

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

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

the class InMemoryReaderTest method testDynamicSplit.

@Test
public void testDynamicSplit() throws Exception {
    List<Integer> elements = Arrays.asList(33, 44, 55, 66, 77, 88);
    // Should initially read elements at indices: 44@1, 55@2, 66@3, 77@4
    Coder<Integer> coder = BigEndianIntegerCoder.of();
    InMemoryReader<Integer> inMemoryReader = new InMemoryReader<>(encodedElements(elements, coder), 1, 4, coder);
    // Unstarted iterator.
    try (InMemoryReader<Integer>.InMemoryReaderIterator iterator = inMemoryReader.iterator()) {
        assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(3L)));
    }
    // Illegal proposed split position.
    try (InMemoryReader<Integer>.InMemoryReaderIterator iterator = inMemoryReader.iterator()) {
        assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(3L)));
        // Poke the iterator so that we can test dynamic splitting.
        assertTrue(iterator.start());
        assertNull(iterator.requestDynamicSplit(toDynamicSplitRequest(new ApproximateSplitRequest())));
        assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(null)));
    }
    // Successful update.
    try (InMemoryReader<Integer>.InMemoryReaderIterator iterator = inMemoryReader.iterator()) {
        // Poke the iterator so that we can test dynamic splitting.
        assertTrue(iterator.start());
        NativeReader.DynamicSplitResult dynamicSplitResult = iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(3L));
        Assert.assertEquals(ReaderTestUtils.positionAtIndex(3L), ReaderTestUtils.positionFromSplitResult(dynamicSplitResult));
        assertEquals(3, iterator.tracker.getStopPosition().longValue());
        assertEquals(44, iterator.getCurrent().intValue());
        assertTrue(iterator.advance());
        assertEquals(55, iterator.getCurrent().intValue());
        assertFalse(iterator.advance());
    }
    // Proposed split position is before the current position, no update.
    try (InMemoryReader<Integer>.InMemoryReaderIterator iterator = inMemoryReader.iterator()) {
        // Poke the iterator so that we can test dynamic splitting.
        assertTrue(iterator.start());
        assertEquals(44, iterator.getCurrent().intValue());
        assertTrue(iterator.advance());
        assertEquals(55, iterator.getCurrent().intValue());
        // Returns true => we promised to return 66.
        assertTrue(iterator.advance());
        // Now we have to refuse the split.
        assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(3L)));
        assertEquals(4, iterator.tracker.getStopPosition().longValue());
        assertEquals(66, iterator.getCurrent().intValue());
        assertFalse(iterator.advance());
    }
    // Proposed split position is after the current stop (end) position, no update.
    try (InMemoryReader<Integer>.InMemoryReaderIterator iterator = inMemoryReader.iterator()) {
        // Poke the iterator so that we can test dynamic splitting.
        assertTrue(iterator.start());
        assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtIndex(5L)));
        assertEquals(4, iterator.tracker.getStopPosition().longValue());
    }
}