From 3f69fa3e9473b51a42bca4d50ff34608ca6cd1b0 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:25:26 +0000 Subject: [PATCH 1/9] test(datastore): use GAX retries in flaky ITDatastoreProtoClientTest --- .../datastore-v1-proto-client/pom.xml | 5 ++ .../client/it/ITDatastoreProtoClientTest.java | 71 +++++++++++++++++-- .../google-cloud-datastore-utils/pom.xml | 5 ++ .../utils/it/ITDatastoreProtoClientTest.java | 71 +++++++++++++++++-- 4 files changed, 144 insertions(+), 8 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/pom.xml b/java-datastore/datastore-v1-proto-client/pom.xml index 3bea6a35fec2..856d5dd0548b 100644 --- a/java-datastore/datastore-v1-proto-client/pom.xml +++ b/java-datastore/datastore-v1-proto-client/pom.xml @@ -107,6 +107,11 @@ + + com.google.api + gax + test + diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 3e6cf025e7f3..0779e314686f 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -27,6 +27,17 @@ import com.google.datastore.v1.client.Datastore; import com.google.datastore.v1.client.DatastoreException; import com.google.datastore.v1.client.DatastoreHelper; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.retrying.DirectRetryingExecutor; +import com.google.api.gax.retrying.RetryAlgorithm; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiClock; +import com.google.api.core.NanoClock; +import java.time.Duration; +import java.util.concurrent.ExecutionException; import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; @@ -59,8 +70,7 @@ public void testQuerySplitterWithDefaultDb() throws DatastoreException { PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).build(); - List splits = - DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE); + List splits = getSplitsWithRetry(query, PARTITION, 2, DATASTORE); Truth.assertThat(splits).isNotEmpty(); splits.forEach( split -> { @@ -81,8 +91,7 @@ public void testQuerySplitterWithDb() throws DatastoreException { PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-db").build(); - List splits = - DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE); + List splits = getSplitsWithRetry(query, PARTITION, 2, DATASTORE); Truth.assertThat(splits).isNotEmpty(); splits.forEach( @@ -91,4 +100,58 @@ public void testQuerySplitterWithDb() throws DatastoreException { Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter); }); } + + // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic + // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry + // here to handle transient backend errors (such as Code.INTERNAL auth issues). + // We reuse GAX retrying utilities here in the test to implement this backoff/retry. + private static List getSplitsWithRetry( + Query query, PartitionId partition, int numSplits, Datastore datastore) + throws DatastoreException { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofSeconds(1)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelayDuration(Duration.ofSeconds(3)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + ApiClock clock = NanoClock.getDefaultClock(); + RetryAlgorithm> retryAlgorithm = new RetryAlgorithm<>( + new BasicResultRetryAlgorithm>() { + @Override + public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + if (prevThrowable instanceof DatastoreException) { + DatastoreException de = (DatastoreException) prevThrowable; + return de.getCode() == com.google.rpc.Code.INTERNAL; + } + return false; + } + }, + new ExponentialRetryAlgorithm(retrySettings, clock) + ); + + DirectRetryingExecutor> executor = new DirectRetryingExecutor<>(retryAlgorithm); + RetryingFuture> future = executor.createFuture( + () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore) + ); + + ApiFuture> submittedFuture = executor.submit(future); + + try { + return submittedFuture.get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof DatastoreException) { + throw (DatastoreException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } } diff --git a/java-datastore/google-cloud-datastore-utils/pom.xml b/java-datastore/google-cloud-datastore-utils/pom.xml index b0533059163a..e5de55f6139e 100644 --- a/java-datastore/google-cloud-datastore-utils/pom.xml +++ b/java-datastore/google-cloud-datastore-utils/pom.xml @@ -84,6 +84,11 @@ + + com.google.api + gax + test + diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index d30c1cbdc598..a7c2c10a344b 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -23,6 +23,17 @@ import com.google.datastore.utils.DatastoreException; import com.google.datastore.utils.DatastoreHelper; import com.google.datastore.v1.*; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.retrying.DirectRetryingExecutor; +import com.google.api.gax.retrying.RetryAlgorithm; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiClock; +import com.google.api.core.NanoClock; +import java.time.Duration; +import java.util.concurrent.ExecutionException; import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; @@ -55,8 +66,7 @@ public void testQuerySplitterWithDefaultDb() throws DatastoreException { PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).build(); - List splits = - DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE); + List splits = getSplitsWithRetry(query, PARTITION, 2, DATASTORE); Truth.assertThat(splits).isNotEmpty(); splits.forEach( split -> { @@ -77,8 +87,7 @@ public void testQuerySplitterWithDb() throws DatastoreException { PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-db").build(); - List splits = - DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE); + List splits = getSplitsWithRetry(query, PARTITION, 2, DATASTORE); Truth.assertThat(splits).isNotEmpty(); splits.forEach( @@ -87,4 +96,58 @@ public void testQuerySplitterWithDb() throws DatastoreException { Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter); }); } + + // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic + // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry + // here to handle transient backend errors (such as Code.INTERNAL auth issues). + // We reuse GAX retrying utilities here in the test to implement this backoff/retry. + private static List getSplitsWithRetry( + Query query, PartitionId partition, int numSplits, Datastore datastore) + throws DatastoreException { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofSeconds(1)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelayDuration(Duration.ofSeconds(3)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + ApiClock clock = NanoClock.getDefaultClock(); + RetryAlgorithm> retryAlgorithm = new RetryAlgorithm<>( + new BasicResultRetryAlgorithm>() { + @Override + public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + if (prevThrowable instanceof DatastoreException) { + DatastoreException de = (DatastoreException) prevThrowable; + return de.getCode() == com.google.rpc.Code.INTERNAL; + } + return false; + } + }, + new ExponentialRetryAlgorithm(retrySettings, clock) + ); + + DirectRetryingExecutor> executor = new DirectRetryingExecutor<>(retryAlgorithm); + RetryingFuture> future = executor.createFuture( + () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore) + ); + + ApiFuture> submittedFuture = executor.submit(future); + + try { + return submittedFuture.get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof DatastoreException) { + throw (DatastoreException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } } From 0ae37c639e74f62f09f4bc134f70f916c50364b1 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:32:03 +0000 Subject: [PATCH 2/9] refactor(datastore): generalize retry logic in ITDatastoreProtoClientTest --- .../client/it/ITDatastoreProtoClientTest.java | 55 +++++++++++-------- .../utils/it/ITDatastoreProtoClientTest.java | 55 +++++++++++-------- 2 files changed, 66 insertions(+), 44 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 0779e314686f..ba479877ebad 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -101,26 +101,19 @@ public void testQuerySplitterWithDb() throws DatastoreException { }); } - // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic - // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry - // here to handle transient backend errors (such as Code.INTERNAL auth issues). - // We reuse GAX retrying utilities here in the test to implement this backoff/retry. - private static List getSplitsWithRetry( - Query query, PartitionId partition, int numSplits, Datastore datastore) - throws DatastoreException { - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofSeconds(1)) - .setRetryDelayMultiplier(2.0) - .setMaxRetryDelayDuration(Duration.ofSeconds(3)) - .setTotalTimeoutDuration(Duration.ofSeconds(10)) - .build(); + @FunctionalInterface + private interface DatastoreCallable { + V call() throws DatastoreException; + } + private static V runWithRetry( + DatastoreCallable callable, RetrySettings retrySettings) + throws DatastoreException { ApiClock clock = NanoClock.getDefaultClock(); - RetryAlgorithm> retryAlgorithm = new RetryAlgorithm<>( - new BasicResultRetryAlgorithm>() { + RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( + new BasicResultRetryAlgorithm() { @Override - public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + public boolean shouldRetry(Throwable prevThrowable, V prevResult) { if (prevThrowable instanceof DatastoreException) { DatastoreException de = (DatastoreException) prevThrowable; return de.getCode() == com.google.rpc.Code.INTERNAL; @@ -131,12 +124,10 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { new ExponentialRetryAlgorithm(retrySettings, clock) ); - DirectRetryingExecutor> executor = new DirectRetryingExecutor<>(retryAlgorithm); - RetryingFuture> future = executor.createFuture( - () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore) - ); + DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); + RetryingFuture future = executor.createFuture(callable::call); - ApiFuture> submittedFuture = executor.submit(future); + ApiFuture submittedFuture = executor.submit(future); try { return submittedFuture.get(); @@ -154,4 +145,24 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { throw new RuntimeException(e); } } + + // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic + // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry + // here to handle transient backend errors (such as Code.INTERNAL auth issues). + // We reuse GAX retrying utilities here in the test to implement this backoff/retry. + private static List getSplitsWithRetry( + Query query, PartitionId partition, int numSplits, Datastore datastore) + throws DatastoreException { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofSeconds(1)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelayDuration(Duration.ofSeconds(3)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + return runWithRetry( + () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), + retrySettings + ); + } } diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index a7c2c10a344b..886da51d4e71 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -97,26 +97,19 @@ public void testQuerySplitterWithDb() throws DatastoreException { }); } - // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic - // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry - // here to handle transient backend errors (such as Code.INTERNAL auth issues). - // We reuse GAX retrying utilities here in the test to implement this backoff/retry. - private static List getSplitsWithRetry( - Query query, PartitionId partition, int numSplits, Datastore datastore) - throws DatastoreException { - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofSeconds(1)) - .setRetryDelayMultiplier(2.0) - .setMaxRetryDelayDuration(Duration.ofSeconds(3)) - .setTotalTimeoutDuration(Duration.ofSeconds(10)) - .build(); + @FunctionalInterface + private interface DatastoreCallable { + V call() throws DatastoreException; + } + private static V runWithRetry( + DatastoreCallable callable, RetrySettings retrySettings) + throws DatastoreException { ApiClock clock = NanoClock.getDefaultClock(); - RetryAlgorithm> retryAlgorithm = new RetryAlgorithm<>( - new BasicResultRetryAlgorithm>() { + RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( + new BasicResultRetryAlgorithm() { @Override - public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + public boolean shouldRetry(Throwable prevThrowable, V prevResult) { if (prevThrowable instanceof DatastoreException) { DatastoreException de = (DatastoreException) prevThrowable; return de.getCode() == com.google.rpc.Code.INTERNAL; @@ -127,12 +120,10 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { new ExponentialRetryAlgorithm(retrySettings, clock) ); - DirectRetryingExecutor> executor = new DirectRetryingExecutor<>(retryAlgorithm); - RetryingFuture> future = executor.createFuture( - () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore) - ); + DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); + RetryingFuture future = executor.createFuture(callable::call); - ApiFuture> submittedFuture = executor.submit(future); + ApiFuture submittedFuture = executor.submit(future); try { return submittedFuture.get(); @@ -150,4 +141,24 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { throw new RuntimeException(e); } } + + // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic + // (unlike the high-level google-cloud-datastore gRPC client). We must explicitly retry + // here to handle transient backend errors (such as Code.INTERNAL auth issues). + // We reuse GAX retrying utilities here in the test to implement this backoff/retry. + private static List getSplitsWithRetry( + Query query, PartitionId partition, int numSplits, Datastore datastore) + throws DatastoreException { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofSeconds(1)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelayDuration(Duration.ofSeconds(3)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + return runWithRetry( + () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), + retrySettings + ); + } } From 45ce01e4a38fc2cde089085b1d4441391dc94b30 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:45:12 +0000 Subject: [PATCH 3/9] test(datastore): add javadocs to retry helper --- .../client/it/ITDatastoreProtoClientTest.java | 19 +++++++++++++++++++ .../utils/it/ITDatastoreProtoClientTest.java | 19 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index ba479877ebad..cf6ffdba3c99 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -101,11 +101,30 @@ public void testQuerySplitterWithDb() throws DatastoreException { }); } + /** + * A functional interface similar to {@link java.util.concurrent.Callable} but specialized + * to throw {@link DatastoreException}. This ensures type safety and avoids having to + * handle generic {@link Exception} in the retry helper. + */ @FunctionalInterface private interface DatastoreCallable { V call() throws DatastoreException; } + /** + * A generic helper method that executes a {@link DatastoreCallable} with retries using the GAX + * retrying framework. + * + *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} + * and a {@link BasicResultRetryAlgorithm} that retries only on transient {@code Code.INTERNAL} + * {@link DatastoreException}s. + * + * @param callable the action to execute + * @param retrySettings the retry configuration (backoff, max attempts, timeouts) + * @return the result of the callable execution + * @throws DatastoreException if the execution fails after all retry attempts, or if a + * non-retryable exception is encountered. + */ private static V runWithRetry( DatastoreCallable callable, RetrySettings retrySettings) throws DatastoreException { diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index 886da51d4e71..3eb29de14562 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -97,11 +97,30 @@ public void testQuerySplitterWithDb() throws DatastoreException { }); } + /** + * A functional interface similar to {@link java.util.concurrent.Callable} but specialized + * to throw {@link DatastoreException}. This ensures type safety and avoids having to + * handle generic {@link Exception} in the retry helper. + */ @FunctionalInterface private interface DatastoreCallable { V call() throws DatastoreException; } + /** + * A generic helper method that executes a {@link DatastoreCallable} with retries using the GAX + * retrying framework. + * + *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} + * and a {@link BasicResultRetryAlgorithm} that retries only on transient {@code Code.INTERNAL} + * {@link DatastoreException}s. + * + * @param callable the action to execute + * @param retrySettings the retry configuration (backoff, max attempts, timeouts) + * @return the result of the callable execution + * @throws DatastoreException if the execution fails after all retry attempts, or if a + * non-retryable exception is encountered. + */ private static V runWithRetry( DatastoreCallable callable, RetrySettings retrySettings) throws DatastoreException { From 1130371a49c4601798c1f5f865dae4f2460cb4a7 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:50:31 +0000 Subject: [PATCH 4/9] refactor(datastore): make retry helper customizable and fail fast --- .../client/it/ITDatastoreProtoClientTest.java | 47 ++++++++++++------- .../utils/it/ITDatastoreProtoClientTest.java | 47 ++++++++++++------- 2 files changed, 58 insertions(+), 36 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index cf6ffdba3c99..880c61597795 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -33,6 +33,7 @@ import com.google.api.gax.retrying.ExponentialRetryAlgorithm; import com.google.api.gax.retrying.BasicResultRetryAlgorithm; import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; import com.google.api.core.ApiFuture; import com.google.api.core.ApiClock; import com.google.api.core.NanoClock; @@ -116,30 +117,25 @@ private interface DatastoreCallable { * retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} - * and a {@link BasicResultRetryAlgorithm} that retries only on transient {@code Code.INTERNAL} - * {@link DatastoreException}s. + * and the custom {@link ResultRetryAlgorithmWithContext}. * * @param callable the action to execute * @param retrySettings the retry configuration (backoff, max attempts, timeouts) + * @param resultRetryAlgorithm the algorithm to determine if a failed attempt should be retried * @return the result of the callable execution * @throws DatastoreException if the execution fails after all retry attempts, or if a * non-retryable exception is encountered. */ private static V runWithRetry( - DatastoreCallable callable, RetrySettings retrySettings) + DatastoreCallable callable, + RetrySettings retrySettings, + ResultRetryAlgorithmWithContext resultRetryAlgorithm) throws DatastoreException { ApiClock clock = NanoClock.getDefaultClock(); + // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm + // as required by DirectRetryingExecutor. RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( - new BasicResultRetryAlgorithm() { - @Override - public boolean shouldRetry(Throwable prevThrowable, V prevResult) { - if (prevThrowable instanceof DatastoreException) { - DatastoreException de = (DatastoreException) prevThrowable; - return de.getCode() == com.google.rpc.Code.INTERNAL; - } - return false; - } - }, + resultRetryAlgorithm, new ExponentialRetryAlgorithm(retrySettings, clock) ); @@ -152,9 +148,13 @@ public boolean shouldRetry(Throwable prevThrowable, V prevResult) { return submittedFuture.get(); } catch (ExecutionException e) { Throwable cause = e.getCause(); + // submittedFuture.get() throws ExecutionException wrapping the actual exception. + // We must explicitly check the type of the cause and unwrap it: + // 1. Rethrow DatastoreException to preserve the method signature. if (cause instanceof DatastoreException) { throw (DatastoreException) cause; } + // 2. Rethrow RuntimeException to avoid wrapping it in another redundant RuntimeException. if (cause instanceof RuntimeException) { throw (RuntimeException) cause; } @@ -172,16 +172,27 @@ public boolean shouldRetry(Throwable prevThrowable, V prevResult) { private static List getSplitsWithRetry( Query query, PartitionId partition, int numSplits, Datastore datastore) throws DatastoreException { + // Fail fast configuration to avoid long wait times during test failures RetrySettings retrySettings = RetrySettings.newBuilder() .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofSeconds(1)) - .setRetryDelayMultiplier(2.0) - .setMaxRetryDelayDuration(Duration.ofSeconds(3)) - .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .setInitialRetryDelayDuration(Duration.ofMillis(200)) + .setRetryDelayMultiplier(1.5) + .setMaxRetryDelayDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(2)) .build(); return runWithRetry( () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), - retrySettings + retrySettings, + new BasicResultRetryAlgorithm>() { + @Override + public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + if (prevThrowable instanceof DatastoreException) { + DatastoreException de = (DatastoreException) prevThrowable; + return de.getCode() == com.google.rpc.Code.INTERNAL; + } + return false; + } + } ); } } diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index 3eb29de14562..775663eb0920 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -29,6 +29,7 @@ import com.google.api.gax.retrying.ExponentialRetryAlgorithm; import com.google.api.gax.retrying.BasicResultRetryAlgorithm; import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; import com.google.api.core.ApiFuture; import com.google.api.core.ApiClock; import com.google.api.core.NanoClock; @@ -112,30 +113,25 @@ private interface DatastoreCallable { * retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} - * and a {@link BasicResultRetryAlgorithm} that retries only on transient {@code Code.INTERNAL} - * {@link DatastoreException}s. + * and the custom {@link ResultRetryAlgorithmWithContext}. * * @param callable the action to execute * @param retrySettings the retry configuration (backoff, max attempts, timeouts) + * @param resultRetryAlgorithm the algorithm to determine if a failed attempt should be retried * @return the result of the callable execution * @throws DatastoreException if the execution fails after all retry attempts, or if a * non-retryable exception is encountered. */ private static V runWithRetry( - DatastoreCallable callable, RetrySettings retrySettings) + DatastoreCallable callable, + RetrySettings retrySettings, + ResultRetryAlgorithmWithContext resultRetryAlgorithm) throws DatastoreException { ApiClock clock = NanoClock.getDefaultClock(); + // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm + // as required by DirectRetryingExecutor. RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( - new BasicResultRetryAlgorithm() { - @Override - public boolean shouldRetry(Throwable prevThrowable, V prevResult) { - if (prevThrowable instanceof DatastoreException) { - DatastoreException de = (DatastoreException) prevThrowable; - return de.getCode() == com.google.rpc.Code.INTERNAL; - } - return false; - } - }, + resultRetryAlgorithm, new ExponentialRetryAlgorithm(retrySettings, clock) ); @@ -148,9 +144,13 @@ public boolean shouldRetry(Throwable prevThrowable, V prevResult) { return submittedFuture.get(); } catch (ExecutionException e) { Throwable cause = e.getCause(); + // submittedFuture.get() throws ExecutionException wrapping the actual exception. + // We must explicitly check the type of the cause and unwrap it: + // 1. Rethrow DatastoreException to preserve the method signature. if (cause instanceof DatastoreException) { throw (DatastoreException) cause; } + // 2. Rethrow RuntimeException to avoid wrapping it in another redundant RuntimeException. if (cause instanceof RuntimeException) { throw (RuntimeException) cause; } @@ -168,16 +168,27 @@ public boolean shouldRetry(Throwable prevThrowable, V prevResult) { private static List getSplitsWithRetry( Query query, PartitionId partition, int numSplits, Datastore datastore) throws DatastoreException { + // Fail fast configuration to avoid long wait times during test failures RetrySettings retrySettings = RetrySettings.newBuilder() .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofSeconds(1)) - .setRetryDelayMultiplier(2.0) - .setMaxRetryDelayDuration(Duration.ofSeconds(3)) - .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .setInitialRetryDelayDuration(Duration.ofMillis(200)) + .setRetryDelayMultiplier(1.5) + .setMaxRetryDelayDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(2)) .build(); return runWithRetry( () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), - retrySettings + retrySettings, + new BasicResultRetryAlgorithm>() { + @Override + public boolean shouldRetry(Throwable prevThrowable, List prevResult) { + if (prevThrowable instanceof DatastoreException) { + DatastoreException de = (DatastoreException) prevThrowable; + return de.getCode() == com.google.rpc.Code.INTERNAL; + } + return false; + } + } ); } } From a6753fb308d7d21ff339c3b829bf4f8798622576 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:54:10 +0000 Subject: [PATCH 5/9] refactor(datastore): simplify retry helper by throwing Exception --- .../client/it/ITDatastoreProtoClientTest.java | 53 +++++-------------- .../utils/it/ITDatastoreProtoClientTest.java | 53 +++++-------------- 2 files changed, 24 insertions(+), 82 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 880c61597795..79a5b39412b1 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -37,8 +37,8 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiClock; import com.google.api.core.NanoClock; +import com.google.rpc.Code; import java.time.Duration; -import java.util.concurrent.ExecutionException; import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; @@ -60,7 +60,7 @@ public void setUp() throws GeneralSecurityException, IOException { } @Test - public void testQuerySplitterWithDefaultDb() throws DatastoreException { + public void testQuerySplitterWithDefaultDb() throws Exception { Filter propertyFilter = makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build(); Query query = @@ -81,7 +81,7 @@ public void testQuerySplitterWithDefaultDb() throws DatastoreException { } @Test - public void testQuerySplitterWithDb() throws DatastoreException { + public void testQuerySplitterWithDb() throws Exception { Filter propertyFilter = makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build(); Query query = @@ -103,18 +103,8 @@ public void testQuerySplitterWithDb() throws DatastoreException { } /** - * A functional interface similar to {@link java.util.concurrent.Callable} but specialized - * to throw {@link DatastoreException}. This ensures type safety and avoids having to - * handle generic {@link Exception} in the retry helper. - */ - @FunctionalInterface - private interface DatastoreCallable { - V call() throws DatastoreException; - } - - /** - * A generic helper method that executes a {@link DatastoreCallable} with retries using the GAX - * retrying framework. + * A generic helper method that executes a {@link java.util.concurrent.Callable} with retries using + * the GAX retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} * and the custom {@link ResultRetryAlgorithmWithContext}. @@ -123,14 +113,13 @@ private interface DatastoreCallable { * @param retrySettings the retry configuration (backoff, max attempts, timeouts) * @param resultRetryAlgorithm the algorithm to determine if a failed attempt should be retried * @return the result of the callable execution - * @throws DatastoreException if the execution fails after all retry attempts, or if a - * non-retryable exception is encountered. + * @throws Exception if the execution fails after all retry attempts. */ private static V runWithRetry( - DatastoreCallable callable, + java.util.concurrent.Callable callable, RetrySettings retrySettings, ResultRetryAlgorithmWithContext resultRetryAlgorithm) - throws DatastoreException { + throws Exception { ApiClock clock = NanoClock.getDefaultClock(); // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm // as required by DirectRetryingExecutor. @@ -140,29 +129,11 @@ private static V runWithRetry( ); DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); - RetryingFuture future = executor.createFuture(callable::call); + RetryingFuture future = executor.createFuture(callable); ApiFuture submittedFuture = executor.submit(future); - try { - return submittedFuture.get(); - } catch (ExecutionException e) { - Throwable cause = e.getCause(); - // submittedFuture.get() throws ExecutionException wrapping the actual exception. - // We must explicitly check the type of the cause and unwrap it: - // 1. Rethrow DatastoreException to preserve the method signature. - if (cause instanceof DatastoreException) { - throw (DatastoreException) cause; - } - // 2. Rethrow RuntimeException to avoid wrapping it in another redundant RuntimeException. - if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; - } - throw new RuntimeException(e); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException(e); - } + return submittedFuture.get(); } // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic @@ -171,7 +142,7 @@ private static V runWithRetry( // We reuse GAX retrying utilities here in the test to implement this backoff/retry. private static List getSplitsWithRetry( Query query, PartitionId partition, int numSplits, Datastore datastore) - throws DatastoreException { + throws Exception { // Fail fast configuration to avoid long wait times during test failures RetrySettings retrySettings = RetrySettings.newBuilder() .setMaxAttempts(3) @@ -188,7 +159,7 @@ private static List getSplitsWithRetry( public boolean shouldRetry(Throwable prevThrowable, List prevResult) { if (prevThrowable instanceof DatastoreException) { DatastoreException de = (DatastoreException) prevThrowable; - return de.getCode() == com.google.rpc.Code.INTERNAL; + return de.getCode() == Code.INTERNAL; } return false; } diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index 775663eb0920..9c7b9b9ac2a6 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -33,8 +33,8 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiClock; import com.google.api.core.NanoClock; +import com.google.rpc.Code; import java.time.Duration; -import java.util.concurrent.ExecutionException; import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; @@ -56,7 +56,7 @@ public void setUp() throws GeneralSecurityException, IOException { } @Test - public void testQuerySplitterWithDefaultDb() throws DatastoreException { + public void testQuerySplitterWithDefaultDb() throws Exception { Filter propertyFilter = makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build(); Query query = @@ -77,7 +77,7 @@ public void testQuerySplitterWithDefaultDb() throws DatastoreException { } @Test - public void testQuerySplitterWithDb() throws DatastoreException { + public void testQuerySplitterWithDb() throws Exception { Filter propertyFilter = makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build(); Query query = @@ -99,18 +99,8 @@ public void testQuerySplitterWithDb() throws DatastoreException { } /** - * A functional interface similar to {@link java.util.concurrent.Callable} but specialized - * to throw {@link DatastoreException}. This ensures type safety and avoids having to - * handle generic {@link Exception} in the retry helper. - */ - @FunctionalInterface - private interface DatastoreCallable { - V call() throws DatastoreException; - } - - /** - * A generic helper method that executes a {@link DatastoreCallable} with retries using the GAX - * retrying framework. + * A generic helper method that executes a {@link java.util.concurrent.Callable} with retries using + * the GAX retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} * and the custom {@link ResultRetryAlgorithmWithContext}. @@ -119,14 +109,13 @@ private interface DatastoreCallable { * @param retrySettings the retry configuration (backoff, max attempts, timeouts) * @param resultRetryAlgorithm the algorithm to determine if a failed attempt should be retried * @return the result of the callable execution - * @throws DatastoreException if the execution fails after all retry attempts, or if a - * non-retryable exception is encountered. + * @throws Exception if the execution fails after all retry attempts. */ private static V runWithRetry( - DatastoreCallable callable, + java.util.concurrent.Callable callable, RetrySettings retrySettings, ResultRetryAlgorithmWithContext resultRetryAlgorithm) - throws DatastoreException { + throws Exception { ApiClock clock = NanoClock.getDefaultClock(); // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm // as required by DirectRetryingExecutor. @@ -136,29 +125,11 @@ private static V runWithRetry( ); DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); - RetryingFuture future = executor.createFuture(callable::call); + RetryingFuture future = executor.createFuture(callable); ApiFuture submittedFuture = executor.submit(future); - try { - return submittedFuture.get(); - } catch (ExecutionException e) { - Throwable cause = e.getCause(); - // submittedFuture.get() throws ExecutionException wrapping the actual exception. - // We must explicitly check the type of the cause and unwrap it: - // 1. Rethrow DatastoreException to preserve the method signature. - if (cause instanceof DatastoreException) { - throw (DatastoreException) cause; - } - // 2. Rethrow RuntimeException to avoid wrapping it in another redundant RuntimeException. - if (cause instanceof RuntimeException) { - throw (RuntimeException) cause; - } - throw new RuntimeException(e); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException(e); - } + return submittedFuture.get(); } // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic @@ -167,7 +138,7 @@ private static V runWithRetry( // We reuse GAX retrying utilities here in the test to implement this backoff/retry. private static List getSplitsWithRetry( Query query, PartitionId partition, int numSplits, Datastore datastore) - throws DatastoreException { + throws Exception { // Fail fast configuration to avoid long wait times during test failures RetrySettings retrySettings = RetrySettings.newBuilder() .setMaxAttempts(3) @@ -184,7 +155,7 @@ private static List getSplitsWithRetry( public boolean shouldRetry(Throwable prevThrowable, List prevResult) { if (prevThrowable instanceof DatastoreException) { DatastoreException de = (DatastoreException) prevThrowable; - return de.getCode() == com.google.rpc.Code.INTERNAL; + return de.getCode() == Code.INTERNAL; } return false; } From 399c44414dcc5e13d17aaa19e2af1b9aa8226a71 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 15:56:33 +0000 Subject: [PATCH 6/9] refactor(datastore): avoid fully qualified Callable in ITDatastoreProtoClientTest --- .../datastore/v1/client/it/ITDatastoreProtoClientTest.java | 5 +++-- .../datastore/utils/it/ITDatastoreProtoClientTest.java | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 79a5b39412b1..700f597c5520 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -42,6 +42,7 @@ import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; +import java.util.concurrent.Callable; import org.junit.Before; import org.junit.Test; @@ -103,7 +104,7 @@ public void testQuerySplitterWithDb() throws Exception { } /** - * A generic helper method that executes a {@link java.util.concurrent.Callable} with retries using + * A generic helper method that executes a {@link Callable} with retries using * the GAX retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} @@ -116,7 +117,7 @@ public void testQuerySplitterWithDb() throws Exception { * @throws Exception if the execution fails after all retry attempts. */ private static V runWithRetry( - java.util.concurrent.Callable callable, + Callable callable, RetrySettings retrySettings, ResultRetryAlgorithmWithContext resultRetryAlgorithm) throws Exception { diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index 9c7b9b9ac2a6..186d05a34112 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -38,6 +38,7 @@ import java.io.IOException; import java.security.GeneralSecurityException; import java.util.List; +import java.util.concurrent.Callable; import org.junit.Before; import org.junit.Test; @@ -99,7 +100,7 @@ public void testQuerySplitterWithDb() throws Exception { } /** - * A generic helper method that executes a {@link java.util.concurrent.Callable} with retries using + * A generic helper method that executes a {@link Callable} with retries using * the GAX retrying framework. * *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} @@ -112,7 +113,7 @@ public void testQuerySplitterWithDb() throws Exception { * @throws Exception if the execution fails after all retry attempts. */ private static V runWithRetry( - java.util.concurrent.Callable callable, + Callable callable, RetrySettings retrySettings, ResultRetryAlgorithmWithContext resultRetryAlgorithm) throws Exception { From 181e2e6f161533c7804b10de461ba1d361ba6361 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 16:04:06 +0000 Subject: [PATCH 7/9] refactor(datastore): unwrap ExecutionException and handle InterruptedException in retry helper --- .../client/it/ITDatastoreProtoClientTest.java | 17 ++++++++++++++++- .../utils/it/ITDatastoreProtoClientTest.java | 17 ++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 700f597c5520..c34427edbc98 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -43,6 +43,7 @@ import java.security.GeneralSecurityException; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; import org.junit.Before; import org.junit.Test; @@ -134,7 +135,21 @@ private static V runWithRetry( ApiFuture submittedFuture = executor.submit(future); - return submittedFuture.get(); + try { + return submittedFuture.get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception) { + throw (Exception) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw e; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } } // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index 186d05a34112..dcdf88e9b126 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -39,6 +39,7 @@ import java.security.GeneralSecurityException; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; import org.junit.Before; import org.junit.Test; @@ -130,7 +131,21 @@ private static V runWithRetry( ApiFuture submittedFuture = executor.submit(future); - return submittedFuture.get(); + try { + return submittedFuture.get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception) { + throw (Exception) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw e; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } } // This low-level Datastore client (proto-over-HTTP) does not have built-in retry logic From f475d79983a46551981e27dc39e998081e77509f Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 16:06:20 +0000 Subject: [PATCH 8/9] docs(datastore): add comments explaining exception unwrapping and interruption handling --- .../datastore/v1/client/it/ITDatastoreProtoClientTest.java | 4 ++++ .../google/datastore/utils/it/ITDatastoreProtoClientTest.java | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index c34427edbc98..6942f5e9e9ce 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -139,6 +139,9 @@ private static V runWithRetry( return submittedFuture.get(); } catch (ExecutionException e) { Throwable cause = e.getCause(); + // submittedFuture.get() wraps any exception thrown during execution in an ExecutionException. + // We unwrap and rethrow the actual cause (Exception or Error) directly so that test failures + // report the root cause (e.g., DatastoreException or AssertionError) instead of the wrapper. if (cause instanceof Exception) { throw (Exception) cause; } @@ -147,6 +150,7 @@ private static V runWithRetry( } throw e; } catch (InterruptedException e) { + // Restore the interrupted status before rethrowing, as per Java concurrency best practices. Thread.currentThread().interrupt(); throw e; } diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index dcdf88e9b126..ab42ba3883f9 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -135,6 +135,9 @@ private static V runWithRetry( return submittedFuture.get(); } catch (ExecutionException e) { Throwable cause = e.getCause(); + // submittedFuture.get() wraps any exception thrown during execution in an ExecutionException. + // We unwrap and rethrow the actual cause (Exception or Error) directly so that test failures + // report the root cause (e.g., DatastoreException or AssertionError) instead of the wrapper. if (cause instanceof Exception) { throw (Exception) cause; } @@ -143,6 +146,7 @@ private static V runWithRetry( } throw e; } catch (InterruptedException e) { + // Restore the interrupted status before rethrowing, as per Java concurrency best practices. Thread.currentThread().interrupt(); throw e; } From 943c79322123e53901d63ff48d717043d1f6eb4e Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 8 Jun 2026 16:53:01 +0000 Subject: [PATCH 9/9] style(datastore): format retry helper in ITDatastoreProtoClientTest --- .../client/it/ITDatastoreProtoClientTest.java | 58 +++++++++---------- .../utils/it/ITDatastoreProtoClientTest.java | 58 +++++++++---------- 2 files changed, 56 insertions(+), 60 deletions(-) diff --git a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java index 6942f5e9e9ce..b5bbf33ff9f0 100644 --- a/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java @@ -18,6 +18,16 @@ import static com.google.datastore.v1.client.DatastoreHelper.makeFilter; import static com.google.datastore.v1.client.DatastoreHelper.makeValue; +import com.google.api.core.ApiClock; +import com.google.api.core.ApiFuture; +import com.google.api.core.NanoClock; +import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.DirectRetryingExecutor; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; +import com.google.api.gax.retrying.RetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.retrying.RetryingFuture; import com.google.common.truth.Truth; import com.google.datastore.v1.Filter; import com.google.datastore.v1.KindExpression; @@ -27,20 +37,10 @@ import com.google.datastore.v1.client.Datastore; import com.google.datastore.v1.client.DatastoreException; import com.google.datastore.v1.client.DatastoreHelper; -import com.google.api.gax.retrying.RetrySettings; -import com.google.api.gax.retrying.DirectRetryingExecutor; -import com.google.api.gax.retrying.RetryAlgorithm; -import com.google.api.gax.retrying.ExponentialRetryAlgorithm; -import com.google.api.gax.retrying.BasicResultRetryAlgorithm; -import com.google.api.gax.retrying.RetryingFuture; -import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiClock; -import com.google.api.core.NanoClock; import com.google.rpc.Code; -import java.time.Duration; import java.io.IOException; import java.security.GeneralSecurityException; +import java.time.Duration; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -105,11 +105,11 @@ public void testQuerySplitterWithDb() throws Exception { } /** - * A generic helper method that executes a {@link Callable} with retries using - * the GAX retrying framework. + * A generic helper method that executes a {@link Callable} with retries using the GAX retrying + * framework. * - *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} - * and the custom {@link ResultRetryAlgorithmWithContext}. + *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} and + * the custom {@link ResultRetryAlgorithmWithContext}. * * @param callable the action to execute * @param retrySettings the retry configuration (backoff, max attempts, timeouts) @@ -125,10 +125,9 @@ private static V runWithRetry( ApiClock clock = NanoClock.getDefaultClock(); // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm // as required by DirectRetryingExecutor. - RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( - resultRetryAlgorithm, - new ExponentialRetryAlgorithm(retrySettings, clock) - ); + RetryAlgorithm retryAlgorithm = + new RetryAlgorithm<>( + resultRetryAlgorithm, new ExponentialRetryAlgorithm(retrySettings, clock)); DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); RetryingFuture future = executor.createFuture(callable); @@ -161,16 +160,16 @@ private static V runWithRetry( // here to handle transient backend errors (such as Code.INTERNAL auth issues). // We reuse GAX retrying utilities here in the test to implement this backoff/retry. private static List getSplitsWithRetry( - Query query, PartitionId partition, int numSplits, Datastore datastore) - throws Exception { + Query query, PartitionId partition, int numSplits, Datastore datastore) throws Exception { // Fail fast configuration to avoid long wait times during test failures - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofMillis(200)) - .setRetryDelayMultiplier(1.5) - .setMaxRetryDelayDuration(Duration.ofMillis(500)) - .setTotalTimeoutDuration(Duration.ofSeconds(2)) - .build(); + RetrySettings retrySettings = + RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofMillis(200)) + .setRetryDelayMultiplier(1.5) + .setMaxRetryDelayDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(2)) + .build(); return runWithRetry( () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), retrySettings, @@ -183,7 +182,6 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { } return false; } - } - ); + }); } } diff --git a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java index ab42ba3883f9..3b7776b941f1 100644 --- a/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java +++ b/java-datastore/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java @@ -18,25 +18,25 @@ import static com.google.datastore.utils.DatastoreHelper.makeFilter; import static com.google.datastore.utils.DatastoreHelper.makeValue; +import com.google.api.core.ApiClock; +import com.google.api.core.ApiFuture; +import com.google.api.core.NanoClock; +import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.DirectRetryingExecutor; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; +import com.google.api.gax.retrying.RetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.retrying.RetryingFuture; import com.google.common.truth.Truth; import com.google.datastore.utils.Datastore; import com.google.datastore.utils.DatastoreException; import com.google.datastore.utils.DatastoreHelper; import com.google.datastore.v1.*; -import com.google.api.gax.retrying.RetrySettings; -import com.google.api.gax.retrying.DirectRetryingExecutor; -import com.google.api.gax.retrying.RetryAlgorithm; -import com.google.api.gax.retrying.ExponentialRetryAlgorithm; -import com.google.api.gax.retrying.BasicResultRetryAlgorithm; -import com.google.api.gax.retrying.RetryingFuture; -import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext; -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiClock; -import com.google.api.core.NanoClock; import com.google.rpc.Code; -import java.time.Duration; import java.io.IOException; import java.security.GeneralSecurityException; +import java.time.Duration; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -101,11 +101,11 @@ public void testQuerySplitterWithDb() throws Exception { } /** - * A generic helper method that executes a {@link Callable} with retries using - * the GAX retrying framework. + * A generic helper method that executes a {@link Callable} with retries using the GAX retrying + * framework. * - *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} - * and the custom {@link ResultRetryAlgorithmWithContext}. + *

It configures a {@link DirectRetryingExecutor} with the provided {@link RetrySettings} and + * the custom {@link ResultRetryAlgorithmWithContext}. * * @param callable the action to execute * @param retrySettings the retry configuration (backoff, max attempts, timeouts) @@ -121,10 +121,9 @@ private static V runWithRetry( ApiClock clock = NanoClock.getDefaultClock(); // We must wrap the result algorithm and timed algorithm into a RetryAlgorithm // as required by DirectRetryingExecutor. - RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( - resultRetryAlgorithm, - new ExponentialRetryAlgorithm(retrySettings, clock) - ); + RetryAlgorithm retryAlgorithm = + new RetryAlgorithm<>( + resultRetryAlgorithm, new ExponentialRetryAlgorithm(retrySettings, clock)); DirectRetryingExecutor executor = new DirectRetryingExecutor<>(retryAlgorithm); RetryingFuture future = executor.createFuture(callable); @@ -157,16 +156,16 @@ private static V runWithRetry( // here to handle transient backend errors (such as Code.INTERNAL auth issues). // We reuse GAX retrying utilities here in the test to implement this backoff/retry. private static List getSplitsWithRetry( - Query query, PartitionId partition, int numSplits, Datastore datastore) - throws Exception { + Query query, PartitionId partition, int numSplits, Datastore datastore) throws Exception { // Fail fast configuration to avoid long wait times during test failures - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(3) - .setInitialRetryDelayDuration(Duration.ofMillis(200)) - .setRetryDelayMultiplier(1.5) - .setMaxRetryDelayDuration(Duration.ofMillis(500)) - .setTotalTimeoutDuration(Duration.ofSeconds(2)) - .build(); + RetrySettings retrySettings = + RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofMillis(200)) + .setRetryDelayMultiplier(1.5) + .setMaxRetryDelayDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(2)) + .build(); return runWithRetry( () -> DatastoreHelper.getQuerySplitter().getSplits(query, partition, numSplits, datastore), retrySettings, @@ -179,7 +178,6 @@ public boolean shouldRetry(Throwable prevThrowable, List prevResult) { } return false; } - } - ); + }); } }