DatastoreStorage implements Closeable { public <T, E extends Exception> T runInTransactionWithRetries(TransactionFunction<T, E> f) throws IOException, E { var retrySettings = datastore.getOptions().getRetrySettings(); var backoff = new ExponentialBackOff.Builder() .setInitialIntervalMillis((int) retrySettings.getInitialRetryDelay().toMillis()) .setMaxIntervalMillis((int) retrySettings.getMaxRetryDelay().toMillis()) .setMultiplier(retrySettings.getRetryDelayMultiplier()) .setMaxElapsedTimeMillis((int) retrySettings.getTotalTimeout().toMillis()) .build(); for (int attempt = 1; ; attempt++) { try { return runInTransaction(f); } catch (DatastoreIOException e) { if (!isRetryableTransactionException(e.getCause())) { throw e; } if (retrySettings.getMaxAttempts() > 0 && attempt >= retrySettings.getMaxAttempts()) { throw e; } var sleepMillis = backoff.nextBackOffMillis(); if (sleepMillis == STOP) { throw e; } log.debug("datastore transaction exception, attempt #{}, retrying in {}ms", attempt, sleepMillis, e); sleepMillis(sleepMillis); } } } DatastoreStorage(CheckedDatastore datastore); @VisibleForTesting  DatastoreStorage(CheckedDatastore datastore,
                   Function<CheckedDatastoreTransaction, DatastoreStorageTransaction> storageTransactionFactory,
                   ExecutorService executor, Logger log); @Override void close(); void updateNextNaturalTrigger(WorkflowId workflowId, TriggerInstantSpec triggerSpec); Map<WorkflowId, Workflow> workflows(); Map<WorkflowId, Workflow> workflows(Set<WorkflowId> workflowIds); List<Workflow> workflows(String componentId); T runInTransactionWithRetries(TransactionFunction<T, E> f);  }
@Test public void runInTransactionWithRetriesShouldCallFunction() throws Exception { final DatastoreStorage storage = new DatastoreStorage(datastore, storageTransactionFactory, executor, logger); final CheckedDatastoreTransaction transaction = datastore.newTransaction(); final DatastoreStorageTransaction storageTransaction = spy(new DatastoreStorageTransaction(transaction)); when(storageTransactionFactory.apply(any())).thenReturn(storageTransaction); when(transactionFunction.apply(any())).thenReturn("foo"); String result = storage.runInTransactionWithRetries(transactionFunction); assertThat(result, is("foo")); verify(transactionFunction).apply(storageTransaction); }
@Test public void runInTransactionWithRetriesShouldPropagateUserException() throws Exception { final Exception expectedException = new FooException(); when(transactionFunction.apply(any())).thenThrow(expectedException); try { storage.runInTransactionWithRetries(transactionFunction); fail("Expected exception!"); } catch (FooException e) { assertThat(e, is(expectedException)); } verify(transactionFunction).apply(any()); }
@Test public void runInTransactionWithRetriesShouldHandleRollbackFailure() throws Exception { var transaction = spy(datastore.newTransaction()); doReturn(transaction).when(datastore).newTransaction(); var rollbackFailure = new DatastoreIOException(new DatastoreException(1, "fail", "error")); doThrow(rollbackFailure).when(transaction).rollback(); var cause = new RuntimeException("foobar"); try { storage.runInTransactionWithRetries(tx -> { throw cause; }); } catch (RuntimeException e) { assertThat(e, is(cause)); } verify(transaction).rollback(); }
@Test public void runInTransactionWithRetriesShouldCallFunctionAndRollbackOnPreCommitConflict() throws Exception { var expectedException = new DatastoreIOException(new DatastoreException(1, "", "")); when(transactionFunction.apply(any())).thenThrow(expectedException); var transaction = spy(datastore.newTransaction()); when(datastore.newTransaction()).thenReturn(transaction); try { storage.runInTransactionWithRetries(transactionFunction); fail("Expected exception!"); } catch (DatastoreIOException e) { assertThat(e, is(expectedException)); } verify(transactionFunction).apply(any()); verify(transaction, never()).commit(); verify(transaction).rollback(); }
@Test public void runInTransactionWithRetriesShouldCallFunctionAndRollbackOnCommitFailure() throws Exception { var commitException = new DatastoreIOException(new DatastoreException(1, "", "")); var transaction = spy(datastore.newTransaction()); doReturn(transaction).when(datastore).newTransaction(); when(transactionFunction.apply(any())).thenReturn(""); doThrow(commitException).when(transaction).commit(); try { storage.runInTransactionWithRetries(transactionFunction); fail("Expected exception!"); } catch (DatastoreIOException e) { assertThat(e, is(commitException)); } verify(transactionFunction).apply(any()); verify(transaction).rollback(); }
@Test public void runInTransactionWithRetriesShouldPropagateCommitExceptionIfRollbackFails() throws Exception { var transaction = spy(datastore.newTransaction()); doReturn(transaction).when(datastore).newTransaction(); when(transactionFunction.apply(any())).thenReturn(""); var commitException = new DatastoreIOException(new DatastoreException(1, "", "")); doThrow(commitException).when(transaction).commit(); var rollbackException = new DatastoreIOException(new DatastoreException(2, "", "")); doThrow(rollbackException).when(transaction).rollback(); try { storage.runInTransactionWithRetries(transactionFunction); fail("Expected exception!"); } catch (DatastoreIOException e) { assertThat(e, is(commitException)); } verify(transactionFunction).apply(any()); verify(transaction).rollback(); verify(logger).debug("Exception on rollback", rollbackException); }
@Test public void runInTransactionWithRetriesShouldThrowIfDatastoreNewTransactionFails() throws Exception { var cause = new DatastoreIOException(new DatastoreException(1, "", "")); doThrow(cause).when(datastore).newTransaction(); var exception = assertThrows( DatastoreIOException.class, () -> storage.runInTransactionWithRetries(transactionFunction) ); assertThat(exception, is(cause)); }