Examples with Context org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context used on opensource projects

Example 1 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context in project beam by apache.

the class RetryManagerTest method testNoFailures.

@Test
public void testNoFailures() throws Exception {
    List<Context> contexts = Lists.newArrayList();
    RetryManager<String, RetryManagerTest.Context> retryManager = new RetryManager<>(Duration.millis(1), Duration.millis(1), 5);
    for (int i = 0; i < 5; ++i) {
        Context context = new Context();
        contexts.add(context);
        retryManager.addOperation(c -> {
            ++c.numStarted;
            return ApiFutures.immediateFuture("yes");
        }, cs -> {
            cs.forEach(c -> ++c.numFailed);
            return RetryType.DONT_RETRY;
        }, c -> ++c.numSucceeded, context);
    }
    contexts.forEach(c -> {
        assertEquals(0, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(0, c.numFailed);
    });
    retryManager.run(true);
    contexts.forEach(c -> {
        assertEquals(1, c.numStarted);
        assertEquals(1, c.numSucceeded);
        assertEquals(0, c.numFailed);
    });
}

Example 2 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context in project beam by apache.

the class RetryManagerTest method testHasSucceeded.

@Test
public void testHasSucceeded() throws Exception {
    List<Context> contexts = Lists.newArrayList();
    RetryManager<String, RetryManagerTest.Context> retryManager = new RetryManager<>(Duration.millis(1), Duration.millis(1), 5);
    for (int i = 0; i < 5; ++i) {
        Context context = new Context();
        contexts.add(context);
        retryManager.addOperation(c -> {
            ++c.numStarted;
            return ApiFutures.immediateFuture("yes");
        }, cs -> {
            cs.forEach(c -> ++c.numFailed);
            return RetryType.DONT_RETRY;
        }, c -> ++c.numSucceeded, c -> false, context);
    }
    contexts.forEach(c -> {
        assertEquals(0, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(0, c.numFailed);
    });
    retryManager.run(true);
    contexts.forEach(c -> {
        assertEquals(1, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(1, c.numFailed);
    });
}

Example 3 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context in project beam by apache.

the class RetryManagerTest method testRetryInOrder.

@Test
public void testRetryInOrder() throws Exception {
    Map<String, Context> contexts = Maps.newHashMap();
    Map<String, Integer> expectedStarts = Maps.newHashMap();
    Map<String, Integer> expectedFailures = Maps.newHashMap();
    RetryManager<String, RetryManagerTest.Context> retryManager = new RetryManager<>(Duration.millis(1), Duration.millis(1), 50);
    for (int i = 0; i < 5; ++i) {
        final int index = i;
        String value = "yes " + i;
        Context context = new Context();
        contexts.put(value, context);
        expectedStarts.put(value, i + 2);
        expectedFailures.put(value, i + 1);
        retryManager.addOperation(c -> {
            // fails all subsequent operations.
            if (c.numStarted <= index) {
                ++c.numStarted;
                RuntimeException e = new RuntimeException("foo");
                return ApiFutures.immediateFailedFuture(e);
            } else {
                ++c.numStarted;
                return ApiFutures.immediateFuture(value);
            }
        }, cs -> {
            cs.forEach(c -> ++c.numFailed);
            return RetryType.RETRY_ALL_OPERATIONS;
        }, c -> ++c.numSucceeded, context);
    }
    contexts.values().forEach(c -> {
        assertEquals(0, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(0, c.numFailed);
    });
    retryManager.run(true);
    contexts.entrySet().forEach(e -> {
        assertEquals((int) expectedStarts.get(e.getKey()), e.getValue().numStarted);
        assertEquals(1, e.getValue().numSucceeded);
        assertEquals((int) expectedFailures.get(e.getKey()), e.getValue().numFailed);
    });
}

Example 4 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context in project beam by apache.

the class StorageApiFinalizeWritesDoFn method finishBundle.

@FinishBundle
@SuppressWarnings({ "nullness" })
public void finishBundle(PipelineOptions pipelineOptions) throws Exception {
    DatasetService datasetService = getDatasetService(pipelineOptions);
    for (Map.Entry<String, Collection<String>> entry : commitStreams.entrySet()) {
        final String tableId = entry.getKey();
        final Collection<String> streamNames = entry.getValue();
        final Set<String> alreadyCommittedStreams = Sets.newHashSet();
        RetryManager<BatchCommitWriteStreamsResponse, Context<BatchCommitWriteStreamsResponse>> retryManager = new RetryManager<>(Duration.standardSeconds(1), Duration.standardMinutes(1), 3);
        retryManager.addOperation(c -> {
            Iterable<String> streamsToCommit = Iterables.filter(streamNames, s -> !alreadyCommittedStreams.contains(s));
            batchCommitOperationsSent.inc();
            return datasetService.commitWriteStreams(tableId, streamsToCommit);
        }, contexts -> {
            LOG.error("BatchCommit failed. tableId " + tableId + " streamNames " + streamNames + " error: " + Iterables.getFirst(contexts, null).getError());
            batchCommitOperationsFailed.inc();
            return RetryType.RETRY_ALL_OPERATIONS;
        }, c -> {
            LOG.info("BatchCommit succeeded for tableId " + tableId + " response " + c.getResult());
            batchCommitOperationsSucceeded.inc();
        }, response -> {
            if (!response.hasCommitTime()) {
                for (StorageError storageError : response.getStreamErrorsList()) {
                    if (storageError.getCode() == StorageErrorCode.STREAM_ALREADY_COMMITTED) {
                        // Make sure that we don't retry any streams that are already committed.
                        alreadyCommittedStreams.add(storageError.getEntity());
                    }
                }
                Iterable<String> streamsToCommit = Iterables.filter(streamNames, s -> !alreadyCommittedStreams.contains(s));
                // retry.
                return Iterables.isEmpty(streamsToCommit);
            }
            return true;
        }, new Context<>());
        retryManager.run(true);
    }
}

Example 5 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context in project beam by apache.

the class RetryManagerTest method testDontRetry.

@Test
public void testDontRetry() throws Exception {
    List<Context> contexts = Lists.newArrayList();
    RetryManager<String, RetryManagerTest.Context> retryManager = new RetryManager<>(Duration.millis(1), Duration.millis(1), 50);
    for (int i = 0; i < 5; ++i) {
        Context context = new Context();
        contexts.add(context);
        String value = "yes " + i;
        retryManager.addOperation(c -> {
            if (c.numStarted == 0) {
                ++c.numStarted;
                RuntimeException e = new RuntimeException("foo");
                return ApiFutures.immediateFailedFuture(e);
            } else {
                ++c.numStarted;
                return ApiFutures.immediateFuture(value);
            }
        }, cs -> {
            cs.forEach(c -> ++c.numFailed);
            return RetryType.DONT_RETRY;
        }, c -> ++c.numSucceeded, context);
    }
    contexts.forEach(c -> {
        assertEquals(0, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(0, c.numFailed);
    });
    retryManager.run(true);
    contexts.forEach(c -> {
        assertEquals(1, c.numStarted);
        assertEquals(0, c.numSucceeded);
        assertEquals(1, c.numFailed);
    });
}