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 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/main/java/io/dgraph/DgraphAsyncClient.java b/src/main/java/io/dgraph/DgraphAsyncClient.java index 40ccb3c..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. */ @@ -97,64 +109,64 @@ 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) { + String errmsg = "error while parsing jwt from the response: "; + LOG.error(errmsg, e); + if (throwOnError) { + throw new AuthException(errmsg, e); + } } finally { wlock.unlock(); } diff --git a/src/test/java/io/dgraph/CompletableFuturesTest.java b/src/test/java/io/dgraph/CompletableFuturesTest.java new file mode 100644 index 0000000..1f4f424 --- /dev/null +++ b/src/test/java/io/dgraph/CompletableFuturesTest.java @@ -0,0 +1,191 @@ +/* + * 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 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); + } + + @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"); + } +}