Skip to content

Commit 5f02d1a

Browse files
Mateusz Krawieccopybara-github
authored andcommitted
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: - `createSharedHttpClient(name)`: OkHttp's default threading, cached and shared. - `createSharedHttpClient(name, executor)`: a client whose dispatcher runs on the given executor. - `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 tests covering client caching, executor injection, and that a caller-injected daemon executor runs `Gemini`'s and `ChatCompletionsHttpClient`'s HTTP work on daemon threads. PiperOrigin-RevId: 945640677
1 parent 410ff81 commit 5f02d1a

7 files changed

Lines changed: 403 additions & 28 deletions

File tree

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.internal.http;
18+
19+
import static java.util.concurrent.TimeUnit.SECONDS;
20+
21+
import java.util.Map;
22+
import java.util.concurrent.ConcurrentHashMap;
23+
import java.util.concurrent.ExecutorService;
24+
import java.util.concurrent.SynchronousQueue;
25+
import java.util.concurrent.ThreadFactory;
26+
import java.util.concurrent.ThreadPoolExecutor;
27+
import java.util.concurrent.atomic.AtomicInteger;
28+
import okhttp3.Dispatcher;
29+
import okhttp3.OkHttpClient;
30+
import org.jspecify.annotations.Nullable;
31+
32+
/**
33+
* Creates {@link OkHttpClient}s for the ADK. The default client is cached per name so its
34+
* dispatcher and connection pool are reused; an executor-backed client is uncached.
35+
*/
36+
public final class HttpClientFactory {
37+
38+
private static final Map<String, OkHttpClient> sharedClients = new ConcurrentHashMap<>();
39+
40+
private HttpClientFactory() {}
41+
42+
/**
43+
* Returns a shared {@link OkHttpClient} cached by {@code name}, using OkHttp's default threading.
44+
* Repeated calls with the same name return the same instance, so the dispatcher and connection
45+
* pool are reused across the ADK.
46+
*
47+
* @param name cache key for the shared client.
48+
*/
49+
public static OkHttpClient createSharedHttpClient(String name) {
50+
return sharedClients.computeIfAbsent(name, unused -> new OkHttpClient());
51+
}
52+
53+
/**
54+
* Returns an {@link OkHttpClient} whose dispatcher runs on {@code executorService}, or the shared
55+
* cached client from {@link #createSharedHttpClient(String)} when it is null. Pass {@link
56+
* #daemonExecutor} so a standalone JVM can exit once work is done, or a container-managed
57+
* executor in a managed environment. The executor-backed client is not cached; the caller owns
58+
* the executor.
59+
*
60+
* @param name cache key used only when {@code executorService} is null.
61+
* @param executorService executor for the dispatcher, or null for the shared default client.
62+
*/
63+
public static OkHttpClient createSharedHttpClient(
64+
String name, @Nullable ExecutorService executorService) {
65+
if (executorService == null) {
66+
return createSharedHttpClient(name);
67+
}
68+
return new OkHttpClient.Builder().dispatcher(new Dispatcher(executorService)).build();
69+
}
70+
71+
/**
72+
* Returns an unbounded pool of daemon threads, matching OkHttp's own dispatcher pool but with
73+
* daemon threads so a standalone JVM can exit once work is done. Managed container environments
74+
* should inject their own executor instead of calling this.
75+
*
76+
* @param name prefix for the dispatcher thread names.
77+
*/
78+
public static ExecutorService daemonExecutor(String name) {
79+
return new ThreadPoolExecutor(
80+
0, Integer.MAX_VALUE, 60L, SECONDS, new SynchronousQueue<>(), daemonThreadFactory(name));
81+
}
82+
83+
private static ThreadFactory daemonThreadFactory(String name) {
84+
AtomicInteger count = new AtomicInteger();
85+
return runnable -> {
86+
Thread thread = new Thread(runnable, name + "-" + count.incrementAndGet());
87+
thread.setDaemon(true);
88+
return thread;
89+
};
90+
}
91+
}

core/src/main/java/com/google/adk/models/Gemini.java

Lines changed: 67 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
2020

2121
import com.google.adk.Version;
22+
import com.google.adk.internal.http.HttpClientFactory;
2223
import com.google.common.collect.ImmutableList;
2324
import com.google.common.collect.ImmutableMap;
2425
import com.google.errorprone.annotations.CanIgnoreReturnValue;
2526
import com.google.genai.Client;
2627
import com.google.genai.ResponseStream;
2728
import com.google.genai.types.Candidate;
29+
import com.google.genai.types.ClientOptions;
2830
import com.google.genai.types.Content;
2931
import com.google.genai.types.FinishReason;
3032
import com.google.genai.types.FunctionCall;
@@ -35,13 +37,16 @@
3537
import com.google.genai.types.Part;
3638
import com.google.genai.types.PartialArg;
3739
import io.reactivex.rxjava3.core.Flowable;
40+
import java.time.Duration;
3841
import java.util.ArrayList;
3942
import java.util.LinkedHashMap;
4043
import java.util.List;
4144
import java.util.Map;
4245
import java.util.Objects;
4346
import java.util.UUID;
4447
import java.util.concurrent.CompletableFuture;
48+
import java.util.concurrent.ExecutorService;
49+
import okhttp3.OkHttpClient;
4550
import org.jspecify.annotations.Nullable;
4651
import org.slf4j.Logger;
4752
import org.slf4j.LoggerFactory;
@@ -57,6 +62,48 @@ public class Gemini extends BaseLlm {
5762
private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
5863
private static final ImmutableMap<String, String> TRACKING_HEADERS;
5964

65+
private static OkHttpClient getSharedHttpClient(@Nullable ExecutorService executorService) {
66+
return HttpClientFactory.createSharedHttpClient("GeminiApiClient", executorService)
67+
.newBuilder()
68+
.connectTimeout(Duration.ZERO)
69+
.readTimeout(Duration.ZERO)
70+
.writeTimeout(Duration.ZERO)
71+
.build();
72+
}
73+
74+
private static Client buildApiKeyClient(
75+
String apiKey, @Nullable ExecutorService executorService) {
76+
return Client.builder()
77+
.apiKey(apiKey)
78+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
79+
.clientOptions(
80+
ClientOptions.builder().customHttpClient(getSharedHttpClient(executorService)).build())
81+
.build();
82+
}
83+
84+
private static Client buildVertexClient(
85+
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
86+
Client.Builder apiClientBuilder =
87+
Client.builder()
88+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
89+
.clientOptions(
90+
ClientOptions.builder()
91+
.customHttpClient(getSharedHttpClient(executorService))
92+
.build());
93+
vertexCredentials.project().ifPresent(apiClientBuilder::project);
94+
vertexCredentials.location().ifPresent(apiClientBuilder::location);
95+
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
96+
return apiClientBuilder.build();
97+
}
98+
99+
private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
100+
return Client.builder()
101+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
102+
.clientOptions(
103+
ClientOptions.builder().customHttpClient(getSharedHttpClient(executorService)).build())
104+
.build();
105+
}
106+
60107
static {
61108
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
62109
String languageLabel = "gl-java/" + JAVA_VERSION.value();
@@ -90,11 +137,7 @@ public Gemini(String modelName, Client apiClient) {
90137
public Gemini(String modelName, String apiKey) {
91138
super(modelName);
92139
Objects.requireNonNull(apiKey, "apiKey cannot be null");
93-
this.apiClient =
94-
Client.builder()
95-
.apiKey(apiKey)
96-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
97-
.build();
140+
this.apiClient = buildApiKeyClient(apiKey, null);
98141
}
99142

100143
/**
@@ -106,12 +149,7 @@ public Gemini(String modelName, String apiKey) {
106149
public Gemini(String modelName, VertexCredentials vertexCredentials) {
107150
super(modelName);
108151
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
109-
Client.Builder apiClientBuilder =
110-
Client.builder().httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build());
111-
vertexCredentials.project().ifPresent(apiClientBuilder::project);
112-
vertexCredentials.location().ifPresent(apiClientBuilder::location);
113-
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
114-
this.apiClient = apiClientBuilder.build();
152+
this.apiClient = buildVertexClient(vertexCredentials, null);
115153
}
116154

117155
/**
@@ -131,6 +169,7 @@ public static class Builder {
131169
private Client apiClient;
132170
private String apiKey;
133171
private VertexCredentials vertexCredentials;
172+
private ExecutorService httpExecutorService;
134173

135174
private Builder() {}
136175

@@ -186,6 +225,20 @@ public Builder vertexCredentials(VertexCredentials vertexCredentials) {
186225
return this;
187226
}
188227

228+
/**
229+
* Sets the executor for the shared HTTP client's dispatcher. Pass {@link
230+
* HttpClientFactory#daemonExecutor} so a standalone or CLI JVM can exit once work is done, or a
231+
* container-managed executor in a managed environment.
232+
*
233+
* @param httpExecutorService The executor for HTTP dispatcher threads.
234+
* @return This builder.
235+
*/
236+
@CanIgnoreReturnValue
237+
public Builder httpExecutorService(ExecutorService httpExecutorService) {
238+
this.httpExecutorService = httpExecutorService;
239+
return this;
240+
}
241+
189242
/**
190243
* Builds the {@link Gemini} instance.
191244
*
@@ -198,15 +251,11 @@ public Gemini build() {
198251
if (apiClient != null) {
199252
return new Gemini(modelName, apiClient);
200253
} else if (apiKey != null) {
201-
return new Gemini(modelName, apiKey);
254+
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService));
202255
} else if (vertexCredentials != null) {
203-
return new Gemini(modelName, vertexCredentials);
256+
return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService));
204257
} else {
205-
return new Gemini(
206-
modelName,
207-
Client.builder()
208-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
209-
.build());
258+
return new Gemini(modelName, buildDefaultClient(httpExecutorService));
210259
}
211260
}
212261
}

core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.fasterxml.jackson.core.JsonProcessingException;
2020
import com.fasterxml.jackson.databind.ObjectMapper;
2121
import com.google.adk.JsonBaseModel;
22+
import com.google.adk.internal.http.HttpClientFactory;
2223
import com.google.adk.models.LlmRequest;
2324
import com.google.adk.models.LlmResponse;
2425
import com.google.common.annotations.VisibleForTesting;
@@ -32,6 +33,7 @@
3233
import java.time.Duration;
3334
import java.util.Map;
3435
import java.util.Objects;
36+
import java.util.concurrent.ExecutorService;
3537
import okhttp3.Call;
3638
import okhttp3.Callback;
3739
import okhttp3.HttpUrl;
@@ -42,6 +44,7 @@
4244
import okhttp3.Response;
4345
import okhttp3.ResponseBody;
4446
import okio.BufferedSource;
47+
import org.jspecify.annotations.Nullable;
4548
import org.slf4j.Logger;
4649
import org.slf4j.LoggerFactory;
4750

@@ -68,11 +71,15 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient {
6871
private static final Duration DEFAULT_CALL_TIMEOUT = Duration.ofMinutes(5);
6972

7073
/**
71-
* Shared OkHttpClient instance whose connection pool and thread dispatcher are reused across all
72-
* {@link ChatCompletionsHttpClient} instances. Each instance forks this client via {@link
73-
* OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking pools.
74+
* Returns the OkHttpClient whose connection pool and thread dispatcher are reused across {@link
75+
* ChatCompletionsHttpClient} instances. When {@code executorService} is null the shared default
76+
* client is returned; otherwise the dispatcher runs on the given executor. Each instance forks
77+
* the result via {@link OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking
78+
* pools.
7479
*/
75-
private static final OkHttpClient SHARED_POOL_CLIENT = new OkHttpClient();
80+
private static OkHttpClient getSharedPoolClient(@Nullable ExecutorService executorService) {
81+
return HttpClientFactory.createSharedHttpClient("ChatCompletionsHttpClient", executorService);
82+
}
7683

7784
private final OkHttpClient client;
7885
private final HttpUrl completionsUrl;
@@ -116,7 +123,19 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient {
116123
* HTTP(S) URL.
117124
*/
118125
public ChatCompletionsHttpClient(HttpOptions httpOptions) {
119-
this(httpOptions, buildClient(httpOptions));
126+
this(httpOptions, buildClient(httpOptions, null));
127+
}
128+
129+
/**
130+
* Constructs a {@link ChatCompletionsHttpClient} whose HTTP dispatcher runs on {@code
131+
* httpExecutorService}. Pass {@link HttpClientFactory#daemonExecutor} so a standalone or CLI JVM
132+
* can exit once work is done, or a container-managed executor in a managed environment.
133+
*
134+
* @param httpOptions HTTP configuration; see {@link #ChatCompletionsHttpClient(HttpOptions)}.
135+
* @param httpExecutorService executor for the HTTP dispatcher threads.
136+
*/
137+
public ChatCompletionsHttpClient(HttpOptions httpOptions, ExecutorService httpExecutorService) {
138+
this(httpOptions, buildClient(httpOptions, httpExecutorService));
120139
}
121140

122141
private ChatCompletionsHttpClient(HttpOptions httpOptions, OkHttpClient client) {
@@ -153,12 +172,13 @@ static ChatCompletionsHttpClient forTesting(HttpOptions httpOptions, OkHttpClien
153172
}
154173

155174
/**
156-
* Builds the production OkHttpClient by forking {@link #SHARED_POOL_CLIENT} so the connection
157-
* pool and dispatcher are reused across instances while applying per-instance timeouts.
175+
* Builds the production OkHttpClient by forking the shared pool client so the connection pool and
176+
* dispatcher are reused across instances while applying per-instance timeouts.
158177
*/
159-
private static OkHttpClient buildClient(HttpOptions httpOptions) {
178+
private static OkHttpClient buildClient(
179+
HttpOptions httpOptions, @Nullable ExecutorService executorService) {
160180
Objects.requireNonNull(httpOptions, "httpOptions cannot be null");
161-
OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder();
181+
OkHttpClient.Builder builder = getSharedPoolClient(executorService).newBuilder();
162182
builder.connectTimeout(Duration.ZERO);
163183
builder.readTimeout(Duration.ZERO);
164184
builder.writeTimeout(Duration.ZERO);

core/src/main/java/com/google/adk/sessions/ApiClient.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
2020

21+
import com.google.adk.internal.http.HttpClientFactory;
2122
import com.google.auth.oauth2.GoogleCredentials;
2223
import com.google.common.base.Ascii;
2324
import com.google.common.base.Strings;
@@ -102,8 +103,12 @@ abstract class ApiClient {
102103
this.httpClient = createHttpClient(httpOptions.timeout().orElse(null));
103104
}
104105

106+
private static OkHttpClient getSharedPoolClient() {
107+
return HttpClientFactory.createSharedHttpClient("ApiClient");
108+
}
109+
105110
private OkHttpClient createHttpClient(@Nullable Integer timeout) {
106-
OkHttpClient.Builder builder = new OkHttpClient().newBuilder();
111+
OkHttpClient.Builder builder = getSharedPoolClient().newBuilder();
107112
if (timeout != null) {
108113
builder.connectTimeout(Duration.ofMillis(timeout));
109114
}

0 commit comments

Comments
 (0)