diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c839f9..e3688a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,41 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.1.0/), ## [Unreleased] +**Changed** + +- `AsyncTransaction.close()` logs a failed abort instead of throwing it, so closing no longer masks + the result of the work the transaction wrapped. Code that catches a close-time failure from an + async transaction no longer sees one. `Transaction.close()` still throws; both javadocs record the + difference. ([#294]) +- `DgraphAsyncClient.withRetry` completes exceptionally with a bare `DgraphException` on every path, + so `whenComplete`, `handle`, and `exceptionally` callbacks receive the `DgraphException` itself. + Retry exhaustion previously handed those callbacks a `CompletionException` wrapping it, while a + first-attempt failure handed them the exception directly. `join()` and `get()` are unaffected: + both wrapped before and still wrap. `DgraphClient.withRetry` runs a separate synchronous retry + loop and is unchanged. ([#294]) + +**Fixed** + +- fix: `DgraphAsyncClient` no longer blocks a `ForkJoinPool.commonPool()` thread for the full + duration of each gRPC call, which could starve the JVM-wide common pool under load. + `CompletableFutures.runWithRetries` now composes on the gRPC future instead of calling a blocking + `.get()`, and `jwt` writes are guarded by the write lock. ([#294]) +- fix: futures returned by `DgraphAsyncClient` complete on the executor given to the constructor, + including the JWT-refresh retry and `withRetry`'s backoff, which previously completed on a gRPC + channel thread and on the common pool respectively. The one exception is a rejection by that + executor, which completes the future on whichever thread observed the rejection. ([#294]) +- fix: a `RejectedExecutionException` from the callback executor no longer leaves the returned + future permanently incomplete. `withRetry` now relays the rejection instead of hanging, and every + rejection surfaces as a `DgraphException` like any other failure. ([#294]) + +**Deprecated** + +- `DgraphAsyncClient(DgraphGrpc.DgraphStub...)` is deprecated in favor of + `DgraphAsyncClient(Executor, DgraphGrpc.DgraphStub...)`. It still defaults to + `ForkJoinPool.commonPool()`, which is unsuitable for I/O continuations: it is a JVM-wide singleton + sized `availableProcessors() - 1`, cannot be tuned per library, and runs unnamed daemon threads + that hide contention in a thread dump. Compiling against it warns; nothing breaks. ([#294]) + ## [25.0.0] - 2026-04-01 **Added** @@ -104,6 +139,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.1.0/), - chore: added a test for best effort queries ([#182]) +[#294]: https://github.com/dgraph-io/dgraph4j/pull/294 [#287]: https://github.com/dgraph-io/dgraph4j/pull/287 [#220]: https://github.com/hypermodeinc/dgraph4j/pull/220 [#215]: https://github.com/hypermodeinc/dgraph4j/pull/215 diff --git a/src/main/java/io/dgraph/AsyncTransaction.java b/src/main/java/io/dgraph/AsyncTransaction.java index 653d23c..209561a 100644 --- a/src/main/java/io/dgraph/AsyncTransaction.java +++ b/src/main/java/io/dgraph/AsyncTransaction.java @@ -12,8 +12,12 @@ import io.dgraph.DgraphProto.TxnContext; import java.util.Collections; import java.util.Map; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.TimeUnit; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * This is the implementation of asynchronous Dgraph transaction. The asynchrony is backed-up by @@ -24,6 +28,7 @@ * @author Michail Klimenkov */ public class AsyncTransaction implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(AsyncTransaction.class); // these can potentially be set from different threads executing the Stub callback private volatile TxnContext context; @@ -390,8 +395,21 @@ private void mergeContext(final TxnContext src) { this.context = builder.build(); } + /** + * Discards the transaction, blocking until the abort completes. Blocks only when the transaction + * has uncommitted mutations, since {@link #discard()} short-circuits otherwise. A failed abort is + * logged rather than thrown: the server cleans up abandoned transactions on its own, and throwing + * here would mask the outcome of the work this transaction wrapped. + * + *

This diverges from {@link Transaction#close()}, which throws. That one is user-invoked + * through try-with-resources, whereas this one also runs from internal cleanup paths. + */ @Override public void close() { - discard().join(); + try { + discard().join(); + } catch (CompletionException | CancellationException e) { + LOG.warn("discarding the transaction during close failed", e); + } } } diff --git a/src/main/java/io/dgraph/CompletableFutures.java b/src/main/java/io/dgraph/CompletableFutures.java index 60c50f3..1676d10 100644 --- a/src/main/java/io/dgraph/CompletableFutures.java +++ b/src/main/java/io/dgraph/CompletableFutures.java @@ -7,6 +7,7 @@ import io.grpc.Context; import java.util.concurrent.*; +import java.util.function.Function; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -40,40 +41,84 @@ static CompletableFuture runWithRetries( Executor executor) { final Callable> ctxCallable = Context.current().wrap(callable); - return CompletableFuture.supplyAsync( - () -> { - try { - return ctxCallable.call().get(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.error("The " + operation + " got interrupted:", e); - throw new DgraphException("The " + operation + " got interrupted", e); - } catch (ExecutionException e) { - if (Exceptions.isJwtExpired(e.getCause())) { - try { - retryLogin.get().get(); - return ctxCallable.call().get(); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - LOG.error("The retried " + operation + " got interrupted:", ie); - throw new DgraphException( - "The retried " + operation + " got interrupted", ie); - } catch (ExecutionException ie) { - LOG.error( - "The retried " + operation + " encounters an execution exception:", ie); - throw new CompletionException(Exceptions.translate(ie.getCause())); - } catch (Exception ie) { - LOG.error( - "The retried " + operation + " encounters a completion exception:", ie); - throw new CompletionException(Exceptions.translate(ie)); - } - } - throw new CompletionException(Exceptions.translate(e.getCause())); - } catch (Exception e) { - throw new CompletionException(Exceptions.translate(e)); - } - }, - executor); + // Composing on the gRPC future rather than blocking on it keeps every stage off the + // critical path: no thread is parked for the round trip. + return invoke(operation, ctxCallable) + .handleAsync( + (value, error) -> + classify(operation, value, error, ctxCallable, retryLogin, executor), + executor) + .thenCompose(Function.identity()) + .handle(CompletableFutures::translateTerminal); + } + + /** + * Upholds the contract that every failure is a {@code CompletionException} wrapping a {@link + * DgraphException}. Runs synchronously so it still fires when {@code executor} rejects a stage. + */ + private static T translateTerminal(T value, Throwable error) { + if (error == null) { + return value; + } + throw new CompletionException(Exceptions.translate(unwrap(error))); + } + + /** Runs the callable, turning a synchronous throw or a null result into a failed future. */ + private static CompletableFuture invoke( + String operation, Callable> callable) { + try { + CompletableFuture future = callable.call(); + if (future != null) { + return future; + } + return CompletableFuture.failedFuture( + new DgraphException("The " + operation + " returned a null future")); + } catch (Exception e) { + return CompletableFuture.failedFuture(e); + } + } + + /** Strips one CompletionException wrapper so error classification sees the real cause. */ + private static Throwable unwrap(Throwable t) { + if (t instanceof CompletionException && t.getCause() != null) { + return t.getCause(); + } + return t; + } + + /** + * Passes a success through, retries once after a JWT refresh on an expired-token failure, and + * translates every other failure to {@code CompletionException(DgraphException)}. All paths + * complete on {@code executor}. + */ + private static CompletableFuture classify( + String operation, + T value, + Throwable error, + Callable> ctxCallable, + Supplier> retryLogin, + Executor executor) { + if (error == null) { + return CompletableFuture.completedFuture(value); + } + + Throwable cause = unwrap(error); + if (Exceptions.isJwtExpired(cause)) { + return retryLogin + .get() + .thenComposeAsync(ignored -> invoke(operation, ctxCallable), executor) + .handleAsync( + (retryValue, retryError) -> { + if (retryError != null) { + LOG.error("The retried {} failed", operation, unwrap(retryError)); + throw new CompletionException(Exceptions.translate(unwrap(retryError))); + } + return retryValue; + }, + executor); + } + + return CompletableFuture.failedFuture(new CompletionException(Exceptions.translate(cause))); } /** @@ -85,13 +130,16 @@ static CompletableFuture runWithRetries( * @param op the operation to execute within a fresh transaction on each attempt * @param attempt the current attempt number (0-based) * @param txnFactory creates a new read-write or read-only transaction per the policy + * @param executor the callback executor; the backoff delay and the returned future's completion + * run on it, and no thread is held for the duration of a delay * @return a CompletableFuture that completes with the result or fails after exhausting retries */ static CompletableFuture attemptAsync( RetryPolicy policy, AsyncTransactionOp op, int attempt, - Supplier txnFactory) { + Supplier txnFactory, + Executor executor) { AsyncTransaction txn = txnFactory.get(); if (policy.isBestEffort()) { @@ -100,8 +148,10 @@ static CompletableFuture attemptAsync( CompletableFuture result = new CompletableFuture<>(); + // handleAsync rather than whenCompleteAsync: the callback consumes the attempt's outcome, so + // the stage below carries only a rejection or a bug in the callback, never a retryable failure. op.execute(txn) - .whenComplete( + .handleAsync( (value, throwable) -> { try { txn.discard(); @@ -111,29 +161,41 @@ static CompletableFuture attemptAsync( if (throwable == null) { result.complete(value); - return; + return null; } DgraphException ex = Exceptions.translate(throwable); if (!ex.isRetryable() || attempt >= policy.getMaxRetries()) { result.completeExceptionally(ex); - return; + return null; } - // Schedule retry after backoff delay long delayMs = policy.calculateDelay(attempt); - Executor delayed = - CompletableFuture.delayedExecutor(delayMs, TimeUnit.MILLISECONDS); - CompletableFuture.supplyAsync(() -> null, delayed) - .thenCompose(ignored -> attemptAsync(policy, op, attempt + 1, txnFactory)) + // The timer takes no executor on purpose: delayedExecutor submits from the internal + // Delayer thread, which swallows a rejection and leaves this future incomplete. + Executor delayed = CompletableFuture.delayedExecutor(delayMs, TimeUnit.MILLISECONDS); + CompletableFuture.runAsync(() -> {}, delayed) + .thenComposeAsync( + ignored -> attemptAsync(policy, op, attempt + 1, txnFactory, executor), + executor) .whenComplete( (retryValue, retryThrowable) -> { if (retryThrowable != null) { - result.completeExceptionally(retryThrowable); + result.completeExceptionally(Exceptions.translate(retryThrowable)); } else { result.complete(retryValue); } }); + return null; + }, + executor) + // A rejection by executor completes this stage exceptionally but leaves result untouched, + // so relay it or the caller waits forever. + .whenComplete( + (ignored, t) -> { + if (t != null) { + result.completeExceptionally(Exceptions.translate(t)); + } }); return result; diff --git a/src/main/java/io/dgraph/DgraphAsyncClient.java b/src/main/java/io/dgraph/DgraphAsyncClient.java index 40ccb3c..ca0c4e1 100644 --- a/src/main/java/io/dgraph/DgraphAsyncClient.java +++ b/src/main/java/io/dgraph/DgraphAsyncClient.java @@ -46,9 +46,18 @@ public class DgraphAsyncClient { * *

A single client is thread safe. * + *

Uses {@link ForkJoinPool#commonPool()} as the callback executor. + * * @param stubs - an array of grpc stubs to be used by this client. The stubs to be used are * chosen at random per transaction. + * @deprecated Use {@link #DgraphAsyncClient(Executor, DgraphGrpc.DgraphStub...)} and supply an + * executor sized for I/O continuations. {@link ForkJoinPool#commonPool()} is a JVM-wide + * singleton sized {@code availableProcessors() - 1}, so a two-vCPU container gets one thread. + * It cannot be tuned per library, its queue is unbounded, and its threads are unnamed + * daemons, which hides contention in a thread dump. This constructor keeps the common pool + * and will be removed in a future major release. */ + @Deprecated public DgraphAsyncClient(DgraphGrpc.DgraphStub... stubs) { this.stubs = asList(stubs); this.executor = ForkJoinPool.commonPool(); @@ -60,7 +69,13 @@ public DgraphAsyncClient(DgraphGrpc.DgraphStub... stubs) { * *

A single client is thread safe. * - * @param executor - the executor to use for various asynchronous tasks executed by this client. + *

The executor is a callback executor: the client runs its continuation logic (JWT + * refresh handling, exception translation, retries) on it, and the futures it returns complete on + * it. gRPC I/O runs on the channel's own threads, and no executor thread is held for the duration + * of a call. Note that the client issues the first attempt of each call on the calling thread, so + * request serialization happens there rather than on the executor. + * + * @param executor the callback executor for this client's continuation logic * @param stubs - an array of grpc stubs to be used by this client. The stubs to be used are * chosen at random per transaction. */ @@ -97,64 +112,62 @@ public CompletableFuture login(String userid, String password) { */ public CompletableFuture loginIntoNamespace( String userid, String password, long namespace) { - Lock wlock = jwtLock.writeLock(); - wlock.lock(); + final DgraphGrpc.DgraphStub client = anyClient(); + final DgraphProto.LoginRequest loginRequest = + DgraphProto.LoginRequest.newBuilder() + .setUserid(userid) + .setPassword(password) + .setNamespace(namespace) + .build(); + + StreamObserverBridge bridge = new StreamObserverBridge<>(); + client.login(loginRequest, bridge); + return bridge.getDelegate().thenAcceptAsync(response -> setJwt(response, true), executor); + } + + protected CompletableFuture retryLogin() { + final String refreshJwt; + Lock rlock = jwtLock.readLock(); + rlock.lock(); try { - final DgraphGrpc.DgraphStub client = anyClient(); - final DgraphProto.LoginRequest loginRequest = - DgraphProto.LoginRequest.newBuilder() - .setUserid(userid) - .setPassword(password) - .setNamespace(namespace) - .build(); - - StreamObserverBridge bridge = new StreamObserverBridge<>(); - client.login(loginRequest, bridge); - return bridge - .getDelegate() - .thenAccept( - (DgraphProto.Response response) -> { - try { - // set the jwt field - jwt = DgraphProto.Jwt.parseFrom(response.getJson()); - } catch (InvalidProtocolBufferException e) { - String errmsg = "error while parsing jwt from the response: "; - LOG.error(errmsg, e); - throw new AuthException(errmsg, e); - } - }); + if (jwt == null || jwt.getRefreshJwt().isEmpty()) { + return CompletableFuture.failedFuture( + new Exception("no refresh JWT available; call login first")); + } + refreshJwt = jwt.getRefreshJwt(); } finally { - wlock.unlock(); + rlock.unlock(); } + + final DgraphGrpc.DgraphStub client = anyClient(); + final DgraphProto.LoginRequest loginRequest = + DgraphProto.LoginRequest.newBuilder().setRefreshToken(refreshJwt).build(); + + StreamObserverBridge bridge = new StreamObserverBridge<>(); + client.login(loginRequest, bridge); + return bridge.getDelegate().thenAcceptAsync(response -> setJwt(response, false), executor); } - protected CompletableFuture retryLogin() { + /** + * Parses the JWT from a login or refresh response and stores it. This is the only writer of the + * {@code jwt} field, and it holds the write lock so the write is published to readers in {@link + * #getStubWithJwt}. + * + * @param response the login or refresh response + * @param throwOnError if true (initial login), a parse failure throws AuthException; if false + * (token refresh), it is logged and swallowed + */ + private void setJwt(DgraphProto.Response response, boolean throwOnError) { Lock wlock = jwtLock.writeLock(); wlock.lock(); try { - if (jwt.getRefreshJwt().isEmpty()) { - CompletableFuture future = new CompletableFuture<>(); - future.completeExceptionally(new Exception("refresh JWT should not be empty")); - return future; + jwt = DgraphProto.Jwt.parseFrom(response.getJson()); + } catch (InvalidProtocolBufferException e) { + String errmsg = "error while parsing jwt from the response: "; + LOG.error(errmsg, e); + if (throwOnError) { + throw new AuthException(errmsg, e); } - - final DgraphGrpc.DgraphStub client = anyClient(); - final DgraphProto.LoginRequest loginRequest = - DgraphProto.LoginRequest.newBuilder().setRefreshToken(jwt.getRefreshJwt()).build(); - - StreamObserverBridge bridge = new StreamObserverBridge<>(); - client.login(loginRequest, bridge); - return bridge - .getDelegate() - .thenAccept( - (DgraphProto.Response response) -> { - try { - // set the jwt field - jwt = DgraphProto.Jwt.parseFrom(response.getJson()); - } catch (InvalidProtocolBufferException e) { - LOG.error("error while parsing jwt from the response: ", e); - } - }); } finally { wlock.unlock(); } @@ -624,11 +637,8 @@ public CompletableFuture withRetry(RetryPolicy policy, AsyncTransactionOp policy, op, 0, - () -> { - AsyncTransaction txn = - policy.isReadOnly() ? newReadOnlyTransaction() : newTransaction(); - return txn; - }); + () -> policy.isReadOnly() ? newReadOnlyTransaction() : newTransaction(), + this.executor); } /** Calls %{@link io.grpc.ManagedChannel#shutdown} on all connections for this client */ diff --git a/src/main/java/io/dgraph/DgraphClient.java b/src/main/java/io/dgraph/DgraphClient.java index 164a87d..0b2dfaa 100644 --- a/src/main/java/io/dgraph/DgraphClient.java +++ b/src/main/java/io/dgraph/DgraphClient.java @@ -372,9 +372,15 @@ public static DgraphGrpc.DgraphStub clientStubFromCloudEndpoint( * *

A single client is thread safe. * + *

Callbacks run on {@link java.util.concurrent.ForkJoinPool#commonPool()}. Prefer {@link + * #DgraphClient(Executor, DgraphGrpc.DgraphStub...)} to isolate this client's callback work. + * * @param stubs - an array of grpc stubs to be used by this client. The stubs to be used are * chosen at random per transaction. */ + // Delegates to the deprecated common-pool constructor on purpose: this overload's contract is + // that it supplies no executor. + @SuppressWarnings("deprecation") public DgraphClient(DgraphGrpc.DgraphStub... stubs) { this.asyncClient = new DgraphAsyncClient(stubs); } diff --git a/src/main/java/io/dgraph/Transaction.java b/src/main/java/io/dgraph/Transaction.java index 0fce3c4..b2d9ee4 100644 --- a/src/main/java/io/dgraph/Transaction.java +++ b/src/main/java/io/dgraph/Transaction.java @@ -235,6 +235,11 @@ public void setBestEffort(boolean bestEffort) { asyncTransaction.setBestEffort(bestEffort); } + /** + * Discards the transaction, throwing if the abort fails. This diverges from {@link + * AsyncTransaction#close()}, which logs the failure instead. Closing here is user-invoked through + * try-with-resources, so the caller can act on the failure. + */ @Override public void close() { Exceptions.withExceptionUnwrapped(this::discard); diff --git a/src/test/java/io/dgraph/CompletableFuturesTest.java b/src/test/java/io/dgraph/CompletableFuturesTest.java new file mode 100644 index 0000000..28bfbb4 --- /dev/null +++ b/src/test/java/io/dgraph/CompletableFuturesTest.java @@ -0,0 +1,414 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.dgraph; + +import static org.testng.Assert.*; + +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import org.testng.annotations.Test; + +public class CompletableFuturesTest { + + private static CompletableFuture failed(Throwable t) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(t); + return f; + } + + private static StatusRuntimeException jwtExpired() { + return Status.UNAUTHENTICATED.withDescription("Token is expired").asRuntimeException(); + } + + private static StatusRuntimeException unavailable() { + return Status.UNAVAILABLE.withDescription("connection refused").asRuntimeException(); + } + + private static final Supplier> NO_RETRY_NEEDED = + () -> CompletableFuture.completedFuture(null); + + /** Waits inside a task, restoring the interrupt flag so shutdownNow still ends the task. */ + private static void awaitQuietly(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + /** Accepts the first {@code limit} submissions, then rejects, like a saturated AbortPolicy. */ + private static final class RejectAfter implements Executor { + private final ExecutorService delegate; + private final int limit; + private final AtomicInteger submissions = new AtomicInteger(); + + RejectAfter(ExecutorService delegate, int limit) { + this.delegate = delegate; + this.limit = limit; + } + + @Override + public void execute(Runnable command) { + if (submissions.incrementAndGet() > limit) { + throw new RejectedExecutionException("saturated"); + } + delegate.execute(command); + } + } + + // Regression test for #293: an in-flight call must not hold an executor thread hostage. + // On the old blocking implementation the lone executor thread parks and the marker never runs. + @Test + public void inFlightCallDoesNotHoldExecutorThread() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + for (int i = 0; i < 3; i++) { + CompletableFuture hung = new CompletableFuture<>(); + CompletableFutures.runWithRetries("op", () -> hung, NO_RETRY_NEEDED, executor); + } + CountDownLatch marker = new CountDownLatch(1); + executor.execute(marker::countDown); + assertTrue( + marker.await(2, TimeUnit.SECONDS), + "executor thread was starved by an in-flight gRPC call"); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void successPassesThrough() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + Callable> callable = + () -> CompletableFuture.completedFuture("ok"); + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, NO_RETRY_NEEDED, executor); + + assertEquals(result.get(2, TimeUnit.SECONDS), "ok"); + } + + @Test + public void nullFutureFromCallableIsTranslated() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + Callable> callable = () -> null; + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, NO_RETRY_NEEDED, executor); + + try { + result.get(2, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof DgraphException, "cause was " + e.getCause()); + } + } + + @Test + public void nonJwtErrorIsTranslatedAndNotRetried() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + AtomicInteger logins = new AtomicInteger(); + Supplier> retryLogin = + () -> { + logins.incrementAndGet(); + return CompletableFuture.completedFuture(null); + }; + Callable> callable = () -> failed(unavailable()); + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, retryLogin, executor); + + try { + result.get(2, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue( + e.getCause() instanceof ConnectionException, "cause was " + e.getCause()); + } + assertEquals(logins.get(), 0, "retryLogin must not run for a non-JWT error"); + } + + @Test + public void jwtExpiryTriggersRetryThenSucceeds() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + AtomicInteger calls = new AtomicInteger(); + AtomicInteger logins = new AtomicInteger(); + Callable> callable = + () -> { + if (calls.incrementAndGet() == 1) { + return failed(jwtExpired()); + } + return CompletableFuture.completedFuture("ok"); + }; + Supplier> retryLogin = + () -> { + logins.incrementAndGet(); + return CompletableFuture.completedFuture(null); + }; + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, retryLogin, executor); + + assertEquals(result.get(2, TimeUnit.SECONDS), "ok"); + assertEquals(calls.get(), 2, "callable should be invoked twice (original + retry)"); + assertEquals(logins.get(), 1, "retryLogin should run exactly once"); + } + + @Test + public void jwtExpiryRetryFailureIsTranslated() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + AtomicInteger calls = new AtomicInteger(); + Callable> callable = + () -> { + if (calls.incrementAndGet() == 1) { + return failed(jwtExpired()); + } + return failed(unavailable()); + }; + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, NO_RETRY_NEEDED, executor); + + try { + result.get(2, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue( + e.getCause() instanceof ConnectionException, "cause was " + e.getCause()); + } + assertEquals(calls.get(), 2); + } + + // The callback executor is the documented completion thread for every path, including the + // JWT-retry path, whose gRPC future completes on a channel thread we do not control. + @Test + public void retryPathCompletesOnCallbackExecutor() throws Exception { + ExecutorService executor = + Executors.newSingleThreadExecutor(r -> new Thread(r, "callback-executor")); + CountDownLatch release = new CountDownLatch(1); + try { + CompletableFuture first = new CompletableFuture<>(); + CompletableFuture retry = new CompletableFuture<>(); + CountDownLatch retryInvoked = new CountDownLatch(1); + AtomicInteger calls = new AtomicInteger(); + Callable> callable = + () -> { + if (calls.incrementAndGet() == 1) { + return first; + } + retryInvoked.countDown(); + return retry; + }; + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, NO_RETRY_NEEDED, executor); + + first.completeExceptionally(jwtExpired()); + assertTrue(retryInvoked.await(5, TimeUnit.SECONDS), "retry was never attempted"); + + // Occupancy, not thread identity: naming the completing thread needs a non-async stage on + // result, and the get() below can drain that stage and run it on the test thread instead. + CountDownLatch occupied = new CountDownLatch(1); + executor.execute( + () -> { + occupied.countDown(); + awaitQuietly(release); + }); + assertTrue(occupied.await(5, TimeUnit.SECONDS), "the executor never ran the blocker"); + + Thread grpcThread = new Thread(() -> retry.complete("ok"), "grpc-thread"); + grpcThread.start(); + grpcThread.join(TimeUnit.SECONDS.toMillis(5)); + assertFalse(result.isDone(), "the retry completed off the callback executor"); + + release.countDown(); + assertEquals(result.get(5, TimeUnit.SECONDS), "ok"); + } finally { + release.countDown(); + executor.shutdownNow(); + } + } + + // attemptAsync's backoff used delayedExecutor(delay, unit), which targets the common pool, and + // completed the result straight from the gRPC thread. Nothing here contacts a server. + @Test + public void retryBackoffAndCompletionRunOnSuppliedExecutor() throws Exception { + ExecutorService executor = + Executors.newSingleThreadExecutor(r -> new Thread(r, "retry-executor")); + ManagedChannel channel = ManagedChannelBuilder.forTarget("localhost:1").usePlaintext().build(); + CountDownLatch release = new CountDownLatch(1); + try { + DgraphAsyncClient client = new DgraphAsyncClient(executor, DgraphGrpc.newStub(channel)); + RetryPolicy policy = + RetryPolicy.builder().maxRetries(2).baseDelay(Duration.ofMillis(10)).jitter(0).build(); + + AtomicInteger attempts = new AtomicInteger(); + List attemptThreads = Collections.synchronizedList(new ArrayList<>()); + CompletableFuture secondAttempt = new CompletableFuture<>(); + CountDownLatch retryStarted = new CountDownLatch(1); + AsyncTransactionOp op = + txn -> { + attemptThreads.add(Thread.currentThread().getName()); + if (attempts.incrementAndGet() == 1) { + return failed(unavailable()); + } + retryStarted.countDown(); + return secondAttempt; + }; + + CompletableFuture result = + CompletableFutures.attemptAsync(policy, op, 0, client::newTransaction, executor); + + assertTrue(retryStarted.await(5, TimeUnit.SECONDS), "the retry never ran"); + assertEquals(attempts.get(), 2, "the retryable failure should be retried once"); + // thenComposeAsync always dispatches through the executor, so this thread name is exact. + assertEquals(attemptThreads.get(1), "retry-executor", "backoff ran off the given executor"); + + // Park the sole executor thread, then complete off-executor: if the completion hop is real + // it queues behind the blocker, and result stays pending until the blocker is released. + CountDownLatch occupied = new CountDownLatch(1); + executor.execute( + () -> { + occupied.countDown(); + awaitQuietly(release); + }); + assertTrue(occupied.await(5, TimeUnit.SECONDS), "the executor never ran the blocker"); + + Thread grpcThread = new Thread(() -> secondAttempt.complete("ok"), "grpc-thread"); + grpcThread.start(); + grpcThread.join(TimeUnit.SECONDS.toMillis(5)); + assertFalse(result.isDone(), "completion ran off the given executor"); + + release.countDown(); + assertEquals(result.get(5, TimeUnit.SECONDS), "ok"); + } finally { + release.countDown(); + channel.shutdownNow(); + executor.shutdownNow(); + } + } + + // A bounded executor with the default AbortPolicy rejects under load. Every rejection must + // surface as a DgraphException, not as a raw RejectedExecutionException. + @Test + public void executorRejectionIsTranslated() throws Exception { + ExecutorService delegate = Executors.newSingleThreadExecutor(); + try { + Executor executor = new RejectAfter(delegate, 0); + Callable> callable = + () -> CompletableFuture.completedFuture("ok"); + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, NO_RETRY_NEEDED, executor); + + try { + result.get(2, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof DgraphException, "cause was " + e.getCause()); + } + } finally { + delegate.shutdownNow(); + } + } + + // attemptAsync completes its result only from inside the callback, so a rejection of that + // callback must be relayed or the caller waits forever. + @Test + public void rejectedCallbackFailsInsteadOfHanging() throws Exception { + ExecutorService delegate = Executors.newSingleThreadExecutor(); + ManagedChannel channel = ManagedChannelBuilder.forTarget("localhost:1").usePlaintext().build(); + try { + Executor executor = new RejectAfter(delegate, 0); + DgraphAsyncClient client = new DgraphAsyncClient(executor, DgraphGrpc.newStub(channel)); + RetryPolicy policy = RetryPolicy.builder().maxRetries(2).build(); + + AsyncTransactionOp op = txn -> CompletableFuture.completedFuture("ok"); + CompletableFuture result = + CompletableFutures.attemptAsync(policy, op, 0, client::newTransaction, executor); + + try { + result.get(5, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof DgraphException, "cause was " + e.getCause()); + } + } finally { + channel.shutdownNow(); + delegate.shutdownNow(); + } + } + + // The backoff hop is the second submission. delayedExecutor must not carry the user executor: + // a rejection raised on the internal Delayer thread is swallowed and the retry never completes. + @Test + public void rejectedBackoffHopFailsInsteadOfHanging() throws Exception { + ExecutorService delegate = Executors.newSingleThreadExecutor(); + ManagedChannel channel = ManagedChannelBuilder.forTarget("localhost:1").usePlaintext().build(); + try { + Executor executor = new RejectAfter(delegate, 1); + DgraphAsyncClient client = new DgraphAsyncClient(executor, DgraphGrpc.newStub(channel)); + RetryPolicy policy = + RetryPolicy.builder().maxRetries(2).baseDelay(Duration.ofMillis(10)).jitter(0).build(); + + AsyncTransactionOp op = txn -> failed(unavailable()); + CompletableFuture result = + CompletableFutures.attemptAsync(policy, op, 0, client::newTransaction, executor); + + try { + result.get(5, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof DgraphException, "cause was " + e.getCause()); + } + } finally { + channel.shutdownNow(); + delegate.shutdownNow(); + } + } + + @Test + public void retryLoginFailureIsTranslated() throws Exception { + Executor executor = ForkJoinPool.commonPool(); + AtomicInteger calls = new AtomicInteger(); + Callable> callable = + () -> { + calls.incrementAndGet(); + return failed(jwtExpired()); + }; + Supplier> retryLogin = + () -> failed(new RuntimeException("refresh failed")); + + CompletableFuture result = + CompletableFutures.runWithRetries("op", callable, retryLogin, executor); + + try { + result.get(2, TimeUnit.SECONDS); + fail("expected failure"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof DgraphException, "cause was " + e.getCause()); + } + assertEquals(calls.get(), 1, "callable must not be retried when login refresh fails"); + } +}