From fbad5261e4f5359f2573f2f717b5eb478cef0dc9 Mon Sep 17 00:00:00 2001 From: Michael Welles Date: Thu, 23 Jul 2026 08:59:46 -0400 Subject: [PATCH 1/4] fix: make DgraphAsyncClient retries non-blocking runWithRetries wrapped each gRPC call in supplyAsync and then blocked the executor thread on .get() for the whole round trip, which starves ForkJoinPool.commonPool() under load. Compose on the stub future instead so no thread is parked; the executor becomes a callback executor. Refs #293 --- .../java/io/dgraph/CompletableFutures.java | 108 +++++++---- .../io/dgraph/CompletableFuturesTest.java | 175 ++++++++++++++++++ 2 files changed, 249 insertions(+), 34 deletions(-) create mode 100644 src/test/java/io/dgraph/CompletableFuturesTest.java diff --git a/src/main/java/io/dgraph/CompletableFutures.java b/src/main/java/io/dgraph/CompletableFutures.java index 60c50f3..a96684b 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,79 @@ 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); + // Fire the RPC (non-blocking) and compose on the resulting future. No thread is + // parked waiting for the round trip; the executor only runs the callback stages. + return invoke(ctxCallable) + .handleAsync( + (value, error) -> + classify(operation, value, error, ctxCallable, retryLogin, executor), + executor) + .thenCompose(Function.identity()); + } + + /** + * Invokes the callable, converting a thrown exception or a null result into an + * already-failed future so the caller never sees a synchronous throw. + */ + private static CompletableFuture invoke(Callable> callable) { + try { + CompletableFuture future = callable.call(); + if (future != null) { + return future; + } + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(new DgraphException("operation returned a null future")); + return failed; + } catch (Exception e) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(e); + return failed; + } + } + + /** 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; + } + + /** + * Classifies the outcome of the first attempt: success passes through; a JWT-expiry + * failure triggers a single login refresh and retry; any other failure is translated. + * Always yields a future that completes with {@code CompletionException(DgraphException)} + * on failure. + */ + 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(ctxCallable), executor) + .handle( + (retryValue, retryError) -> { + if (retryError != null) { + LOG.error("The retried {} failed", operation, unwrap(retryError)); + throw new CompletionException(Exceptions.translate(unwrap(retryError))); + } + return retryValue; + }); + } + + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(new CompletionException(Exceptions.translate(cause))); + return failed; } /** diff --git a/src/test/java/io/dgraph/CompletableFuturesTest.java b/src/test/java/io/dgraph/CompletableFuturesTest.java new file mode 100644 index 0000000..9085e99 --- /dev/null +++ b/src/test/java/io/dgraph/CompletableFuturesTest.java @@ -0,0 +1,175 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.dgraph; + +import static org.testng.Assert.*; + +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +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.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); + + // The regression test for #293: an in-flight (never-completing) call must not + // hold an executor thread hostage. On the old blocking implementation the + // single executor thread parks on .get() 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 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); + } + + @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"); + } +} From 9f9dc20f35f7de29c23f3ca22e822630089ef963 Mon Sep 17 00:00:00 2001 From: Michael Welles Date: Thu, 23 Jul 2026 09:08:22 -0400 Subject: [PATCH 2/4] fix: guard jwt field writes under the write lock The jwt field was written inside a thenAccept callback after the write lock was released, leaving the write unguarded and unpublished across threads. Move the write into a lock-held setter and read the refresh token under the read lock. Refs #293 --- .../java/io/dgraph/DgraphAsyncClient.java | 99 ++++++++++--------- 1 file changed, 50 insertions(+), 49 deletions(-) diff --git a/src/main/java/io/dgraph/DgraphAsyncClient.java b/src/main/java/io/dgraph/DgraphAsyncClient.java index 40ccb3c..e0a1357 100644 --- a/src/main/java/io/dgraph/DgraphAsyncClient.java +++ b/src/main/java/io/dgraph/DgraphAsyncClient.java @@ -97,64 +97,65 @@ public CompletableFuture login(String userid, String password) { */ public CompletableFuture loginIntoNamespace( String userid, String password, long namespace) { - Lock wlock = jwtLock.writeLock(); - wlock.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); - } - }); - } finally { - wlock.unlock(); - } + 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(response -> setJwt(response, true)); } protected CompletableFuture retryLogin() { - Lock wlock = jwtLock.writeLock(); - wlock.lock(); + final String refreshJwt; + Lock rlock = jwtLock.readLock(); + rlock.lock(); try { - if (jwt.getRefreshJwt().isEmpty()) { + if (jwt == null || jwt.getRefreshJwt().isEmpty()) { CompletableFuture future = new CompletableFuture<>(); future.completeExceptionally(new Exception("refresh JWT should not be empty")); return future; } + refreshJwt = jwt.getRefreshJwt(); + } finally { + rlock.unlock(); + } - 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); - } - }); + 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().thenAccept(response -> setJwt(response, false)); + } + + /** + * Parses the JWT from a login/refresh response and stores it under the write lock. This is the + * only writer of the {@code jwt} field; running it under the write lock (rather than around the + * RPC setup, as before) is what makes the write safely published to readers that hold the read + * lock 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, preserving prior behavior + */ + private void setJwt(DgraphProto.Response response, boolean throwOnError) { + Lock wlock = jwtLock.writeLock(); + wlock.lock(); + try { + jwt = DgraphProto.Jwt.parseFrom(response.getJson()); + } catch (InvalidProtocolBufferException e) { + if (throwOnError) { + String errmsg = "error while parsing jwt from the response: "; + LOG.error(errmsg, e); + throw new AuthException(errmsg, e); + } + LOG.error("error while parsing jwt from the response: ", e); } finally { wlock.unlock(); } From f175cb2e03e019918c94c86495d81597a8b6a061 Mon Sep 17 00:00:00 2001 From: Michael Welles Date: Thu, 23 Jul 2026 09:22:19 -0400 Subject: [PATCH 3/4] docs: document callback-executor semantics; address review nits Document in the DgraphAsyncClient constructors that the Executor is a callback executor and that the commonPool default is safe because callbacks never block. Also dedupe a log message in setJwt and cover the null-future path in runWithRetries. Refs #293 --- .../java/io/dgraph/DgraphAsyncClient.java | 19 +++++++++++++++---- .../io/dgraph/CompletableFuturesTest.java | 16 ++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/main/java/io/dgraph/DgraphAsyncClient.java b/src/main/java/io/dgraph/DgraphAsyncClient.java index e0a1357..5231b97 100644 --- a/src/main/java/io/dgraph/DgraphAsyncClient.java +++ b/src/main/java/io/dgraph/DgraphAsyncClient.java @@ -46,6 +46,11 @@ public class DgraphAsyncClient { * *

A single client is thread safe. * + *

Uses {@link ForkJoinPool#commonPool()} as the callback executor. This is safe because the + * client's callbacks never block; use + * {@link #DgraphAsyncClient(Executor, DgraphGrpc.DgraphStub...)} to supply a dedicated executor + * if you want 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. */ @@ -60,7 +65,14 @@ 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 returned futures complete on it. + * gRPC I/O runs on the channel's own threads, and the client never blocks an executor thread for + * the duration of a call. Because these callbacks never block, the no-arg constructor's default + * of {@link ForkJoinPool#commonPool()} is safe; supply your own executor to isolate this + * client's callback work from the common pool. + * + * @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. */ @@ -150,12 +162,11 @@ private void setJwt(DgraphProto.Response response, boolean throwOnError) { try { jwt = DgraphProto.Jwt.parseFrom(response.getJson()); } catch (InvalidProtocolBufferException e) { + String errmsg = "error while parsing jwt from the response: "; + LOG.error(errmsg, e); if (throwOnError) { - String errmsg = "error while parsing jwt from the response: "; - LOG.error(errmsg, e); throw new AuthException(errmsg, e); } - LOG.error("error while parsing jwt from the response: ", e); } finally { wlock.unlock(); } diff --git a/src/test/java/io/dgraph/CompletableFuturesTest.java b/src/test/java/io/dgraph/CompletableFuturesTest.java index 9085e99..1f4f424 100644 --- a/src/test/java/io/dgraph/CompletableFuturesTest.java +++ b/src/test/java/io/dgraph/CompletableFuturesTest.java @@ -74,6 +74,22 @@ public void successPassesThrough() throws Exception { 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(); From a445909a12028fa351156bf7afc9be5d863490bc Mon Sep 17 00:00:00 2001 From: Michael Welles Date: Thu, 23 Jul 2026 21:56:38 -0400 Subject: [PATCH 4/4] chore: add CHANGELOG entry for the non-blocking async client fix Refs #294 --- CHANGELOG.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c839f9..5a2a23c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.1.0/), ## [Unreleased] +**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]) + ## [25.0.0] - 2026-04-01 **Added** @@ -104,6 +111,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