From 2394a9501a15470eba5a164dadf76fa28aeb649b Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Tue, 14 Jul 2026 01:46:01 -0700 Subject: [PATCH] feat: share a single OkHttpClient with injectable daemon threads across the ADK Building a separate `OkHttpClient` per model or session leaks connection pools and threads. This routes `Gemini`, `ChatCompletionsHttpClient`, and the sessions `ApiClient` through `HttpClientFactory`, which caches one shared client per name so the dispatcher and connection pool are reused. Threading is injectable at construction rather than hard-coded: - `getOrCreateSharedHttpClient(name)`: the shared client, OkHttp's default threading. - `createHttpClient(executor)`: a new client whose dispatcher runs on the given executor; not cached, so the caller owns the executor and the returned client. - `daemonExecutor(name)`: an unbounded daemon-thread pool so a standalone or CLI JVM can exit once work is done. - `Gemini.builder().httpExecutorService(...)` and a new `ChatCompletionsHttpClient(HttpOptions, ExecutorService)` constructor accept the executor at construction. The sessions client issues synchronous requests, which do not use the dispatcher executor, so it shares the pool only. `Gemini` forks the shared client with zeroed connect/read/write timeouts to keep the prior infinite-timeout behavior. Adds `HttpClientFactoryTest` covering caching by name, executor injection, and daemon threads. PiperOrigin-RevId: 947525602 --- .../adk/internal/http/HttpClientFactory.java | 82 ++++++++++++++++ .../java/com/google/adk/models/Gemini.java | 91 ++++++++++++++---- .../chat/ChatCompletionsHttpClient.java | 37 ++++++-- .../com/google/adk/sessions/ApiClient.java | 7 +- .../internal/http/HttpClientFactoryTest.java | 93 +++++++++++++++++++ 5 files changed, 283 insertions(+), 27 deletions(-) create mode 100644 core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java create mode 100644 core/src/test/java/com/google/adk/internal/http/HttpClientFactoryTest.java diff --git a/core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java b/core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java new file mode 100644 index 000000000..61a78ce27 --- /dev/null +++ b/core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java @@ -0,0 +1,82 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.internal.http; + +import static java.util.concurrent.TimeUnit.SECONDS; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicInteger; +import okhttp3.Dispatcher; +import okhttp3.OkHttpClient; + +/** + * Creates {@link OkHttpClient}s for the ADK. The default clients are cached per name so the + * dispatcher and connection pool are reused across the ADK; a caller that supplies its own executor + * gets a fresh, non-cached client it owns. + */ +public final class HttpClientFactory { + + private static final Map sharedClients = new ConcurrentHashMap<>(); + + private HttpClientFactory() {} + + /** + * Returns the shared {@link OkHttpClient} cached by {@code name}, using OkHttp's default + * threading. + */ + public static OkHttpClient getOrCreateSharedHttpClient(String name) { + return sharedClients.computeIfAbsent(name, unused -> new OkHttpClient()); + } + + /** + * Returns a new {@link OkHttpClient} whose dispatcher runs on {@code executorService}. Pass + * {@link #daemonExecutor} so a standalone JVM can exit once work is done, or a container-managed + * executor in a managed environment. The client is not cached: the caller owns the executor and + * the returned client. + * + * @param executorService executor for the dispatcher. + */ + public static OkHttpClient createHttpClient(ExecutorService executorService) { + return new OkHttpClient.Builder().dispatcher(new Dispatcher(executorService)).build(); + } + + /** + * Returns an unbounded pool of daemon threads, matching OkHttp's own dispatcher pool but with + * daemon threads so a standalone JVM can exit once work is done. Managed container environments + * should inject their own executor instead of calling this. + * + * @param name prefix for the dispatcher thread names. + */ + public static ExecutorService daemonExecutor(String name) { + return new ThreadPoolExecutor( + 0, Integer.MAX_VALUE, 60L, SECONDS, new SynchronousQueue<>(), daemonThreadFactory(name)); + } + + private static ThreadFactory daemonThreadFactory(String name) { + AtomicInteger count = new AtomicInteger(); + return runnable -> { + Thread thread = new Thread(runnable, name + "-" + count.incrementAndGet()); + thread.setDaemon(true); + return thread; + }; + } +} diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index e13305540..16781e857 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -19,12 +19,14 @@ import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; import com.google.adk.Version; +import com.google.adk.internal.http.HttpClientFactory; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.Client; import com.google.genai.ResponseStream; import com.google.genai.types.Candidate; +import com.google.genai.types.ClientOptions; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; import com.google.genai.types.FunctionCall; @@ -35,6 +37,7 @@ import com.google.genai.types.Part; import com.google.genai.types.PartialArg; import io.reactivex.rxjava3.core.Flowable; +import java.time.Duration; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; @@ -42,6 +45,8 @@ import java.util.Objects; import java.util.UUID; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import okhttp3.OkHttpClient; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -57,6 +62,52 @@ public class Gemini extends BaseLlm { private static final Logger logger = LoggerFactory.getLogger(Gemini.class); private static final ImmutableMap TRACKING_HEADERS; + private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) { + OkHttpClient client = + executorService == null + ? HttpClientFactory.getOrCreateSharedHttpClient("GeminiApiClient") + : HttpClientFactory.createHttpClient(executorService); + return client + .newBuilder() + .connectTimeout(Duration.ZERO) + .readTimeout(Duration.ZERO) + .writeTimeout(Duration.ZERO) + .build(); + } + + private static Client buildApiKeyClient( + String apiKey, @Nullable ExecutorService executorService) { + return Client.builder() + .apiKey(apiKey) + .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) + .clientOptions( + ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()) + .build(); + } + + private static Client buildVertexClient( + VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) { + Client.Builder apiClientBuilder = + Client.builder() + .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) + .clientOptions( + ClientOptions.builder() + .customHttpClient(prepareHttpClient(executorService)) + .build()); + vertexCredentials.project().ifPresent(apiClientBuilder::project); + vertexCredentials.location().ifPresent(apiClientBuilder::location); + vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials); + return apiClientBuilder.build(); + } + + private static Client buildDefaultClient(@Nullable ExecutorService executorService) { + return Client.builder() + .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) + .clientOptions( + ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()) + .build(); + } + static { String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION; String languageLabel = "gl-java/" + JAVA_VERSION.value(); @@ -90,11 +141,7 @@ public Gemini(String modelName, Client apiClient) { public Gemini(String modelName, String apiKey) { super(modelName); Objects.requireNonNull(apiKey, "apiKey cannot be null"); - this.apiClient = - Client.builder() - .apiKey(apiKey) - .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) - .build(); + this.apiClient = buildApiKeyClient(apiKey, null); } /** @@ -106,12 +153,7 @@ public Gemini(String modelName, String apiKey) { public Gemini(String modelName, VertexCredentials vertexCredentials) { super(modelName); Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"); - Client.Builder apiClientBuilder = - Client.builder().httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()); - vertexCredentials.project().ifPresent(apiClientBuilder::project); - vertexCredentials.location().ifPresent(apiClientBuilder::location); - vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials); - this.apiClient = apiClientBuilder.build(); + this.apiClient = buildVertexClient(vertexCredentials, null); } /** @@ -131,6 +173,7 @@ public static class Builder { private Client apiClient; private String apiKey; private VertexCredentials vertexCredentials; + private ExecutorService httpExecutorService; private Builder() {} @@ -186,6 +229,22 @@ public Builder vertexCredentials(VertexCredentials vertexCredentials) { return this; } + /** + * Sets the executor for the shared HTTP client's dispatcher. Pass {@link + * HttpClientFactory#daemonExecutor} so a standalone or CLI JVM can exit once work is done, or a + * container-managed executor in a managed environment. Applies only when the client is built + * from an API key, Vertex credentials, or the default; it is ignored when an explicit {@link + * #apiClient} is supplied. + * + * @param httpExecutorService The executor for HTTP dispatcher threads. + * @return This builder. + */ + @CanIgnoreReturnValue + public Builder httpExecutorService(ExecutorService httpExecutorService) { + this.httpExecutorService = httpExecutorService; + return this; + } + /** * Builds the {@link Gemini} instance. * @@ -198,15 +257,11 @@ public Gemini build() { if (apiClient != null) { return new Gemini(modelName, apiClient); } else if (apiKey != null) { - return new Gemini(modelName, apiKey); + return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService)); } else if (vertexCredentials != null) { - return new Gemini(modelName, vertexCredentials); + return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService)); } else { - return new Gemini( - modelName, - Client.builder() - .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) - .build()); + return new Gemini(modelName, buildDefaultClient(httpExecutorService)); } } } diff --git a/core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java b/core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java index f83287096..236e9b56e 100644 --- a/core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java +++ b/core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java @@ -19,6 +19,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.JsonBaseModel; +import com.google.adk.internal.http.HttpClientFactory; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; import com.google.common.annotations.VisibleForTesting; @@ -32,6 +33,7 @@ import java.time.Duration; import java.util.Map; import java.util.Objects; +import java.util.concurrent.ExecutorService; import okhttp3.Call; import okhttp3.Callback; import okhttp3.HttpUrl; @@ -42,6 +44,7 @@ import okhttp3.Response; import okhttp3.ResponseBody; import okio.BufferedSource; +import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -68,11 +71,16 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient { private static final Duration DEFAULT_CALL_TIMEOUT = Duration.ofMinutes(5); /** - * Shared OkHttpClient instance whose connection pool and thread dispatcher are reused across all - * {@link ChatCompletionsHttpClient} instances. Each instance forks this client via {@link + * Returns the OkHttpClient whose connection pool and thread dispatcher back {@link + * ChatCompletionsHttpClient} instances. Without an executor this is the shared client cached by + * name; with one it is a fresh client the caller owns. Each instance forks it via {@link * OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking pools. */ - private static final OkHttpClient SHARED_POOL_CLIENT = new OkHttpClient(); + private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) { + return executorService == null + ? HttpClientFactory.getOrCreateSharedHttpClient("ChatCompletionsHttpClient") + : HttpClientFactory.createHttpClient(executorService); + } private final OkHttpClient client; private final HttpUrl completionsUrl; @@ -116,7 +124,19 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient { * HTTP(S) URL. */ public ChatCompletionsHttpClient(HttpOptions httpOptions) { - this(httpOptions, buildClient(httpOptions)); + this(httpOptions, buildClient(httpOptions, null)); + } + + /** + * Constructs a {@link ChatCompletionsHttpClient} whose HTTP dispatcher runs on {@code + * httpExecutorService}. Pass {@link HttpClientFactory#daemonExecutor} so a standalone or CLI JVM + * can exit once work is done, or a container-managed executor in a managed environment. + * + * @param httpOptions HTTP configuration; see {@link #ChatCompletionsHttpClient(HttpOptions)}. + * @param httpExecutorService executor for the HTTP dispatcher threads. + */ + public ChatCompletionsHttpClient(HttpOptions httpOptions, ExecutorService httpExecutorService) { + this(httpOptions, buildClient(httpOptions, httpExecutorService)); } private ChatCompletionsHttpClient(HttpOptions httpOptions, OkHttpClient client) { @@ -153,12 +173,13 @@ static ChatCompletionsHttpClient forTesting(HttpOptions httpOptions, OkHttpClien } /** - * Builds the production OkHttpClient by forking {@link #SHARED_POOL_CLIENT} so the connection - * pool and dispatcher are reused across instances while applying per-instance timeouts. + * Builds the production OkHttpClient by forking the shared pool client so the connection pool and + * dispatcher are reused across instances while applying per-instance timeouts. */ - private static OkHttpClient buildClient(HttpOptions httpOptions) { + private static OkHttpClient buildClient( + HttpOptions httpOptions, @Nullable ExecutorService executorService) { Objects.requireNonNull(httpOptions, "httpOptions cannot be null"); - OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder(); + OkHttpClient.Builder builder = prepareHttpClient(executorService).newBuilder(); builder.connectTimeout(Duration.ZERO); builder.readTimeout(Duration.ZERO); builder.writeTimeout(Duration.ZERO); diff --git a/core/src/main/java/com/google/adk/sessions/ApiClient.java b/core/src/main/java/com/google/adk/sessions/ApiClient.java index 1b0485dd2..3b42f6693 100644 --- a/core/src/main/java/com/google/adk/sessions/ApiClient.java +++ b/core/src/main/java/com/google/adk/sessions/ApiClient.java @@ -18,6 +18,7 @@ import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; +import com.google.adk.internal.http.HttpClientFactory; import com.google.auth.oauth2.GoogleCredentials; import com.google.common.base.Ascii; import com.google.common.base.Strings; @@ -102,8 +103,12 @@ abstract class ApiClient { this.httpClient = createHttpClient(httpOptions.timeout().orElse(null)); } + private static OkHttpClient getSharedPoolClient() { + return HttpClientFactory.getOrCreateSharedHttpClient("ApiClient"); + } + private OkHttpClient createHttpClient(@Nullable Integer timeout) { - OkHttpClient.Builder builder = new OkHttpClient().newBuilder(); + OkHttpClient.Builder builder = getSharedPoolClient().newBuilder(); if (timeout != null) { builder.connectTimeout(Duration.ofMillis(timeout)); } diff --git a/core/src/test/java/com/google/adk/internal/http/HttpClientFactoryTest.java b/core/src/test/java/com/google/adk/internal/http/HttpClientFactoryTest.java new file mode 100644 index 000000000..8f146d099 --- /dev/null +++ b/core/src/test/java/com/google/adk/internal/http/HttpClientFactoryTest.java @@ -0,0 +1,93 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.internal.http; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import okhttp3.OkHttpClient; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link HttpClientFactory}. */ +@RunWith(JUnit4.class) +public final class HttpClientFactoryTest { + + @Test + public void getOrCreateSharedHttpClient_sameName_returnsCachedInstance() { + OkHttpClient first = HttpClientFactory.getOrCreateSharedHttpClient("cacheByName"); + OkHttpClient second = HttpClientFactory.getOrCreateSharedHttpClient("cacheByName"); + + assertSame(first, second); + } + + @Test + public void getOrCreateSharedHttpClient_differentNames_returnDistinctInstances() { + OkHttpClient first = HttpClientFactory.getOrCreateSharedHttpClient("nameA"); + OkHttpClient second = HttpClientFactory.getOrCreateSharedHttpClient("nameB"); + + assertNotSame(first, second); + } + + @Test + public void createHttpClient_usesInjectedExecutor() { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + OkHttpClient client = HttpClientFactory.createHttpClient(executor); + + assertSame(executor, client.dispatcher().executorService()); + } finally { + executor.shutdown(); + } + } + + @Test + public void createHttpClient_isNotCached_returnsDistinctInstances() { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + OkHttpClient first = HttpClientFactory.createHttpClient(executor); + OkHttpClient second = HttpClientFactory.createHttpClient(executor); + + assertNotSame(first, second); + } finally { + executor.shutdown(); + } + } + + @Test + public void daemonExecutor_producesDaemonThreads() { + ExecutorService executor = HttpClientFactory.daemonExecutor("daemonPool"); + try { + ThreadPoolExecutor pool = (ThreadPoolExecutor) executor; + assertEquals(0, pool.getCorePoolSize()); + assertEquals(Integer.MAX_VALUE, pool.getMaximumPoolSize()); + assertEquals(60L, pool.getKeepAliveTime(TimeUnit.SECONDS)); + + Thread thread = pool.getThreadFactory().newThread(() -> {}); + assertTrue(thread.isDaemon()); + } finally { + executor.shutdown(); + } + } +}