Skip to content

Commit d3add2d

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: - `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: 945640677
1 parent 6bae658 commit d3add2d

5 files changed

Lines changed: 283 additions & 27 deletions

File tree

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
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+
31+
/**
32+
* Creates {@link OkHttpClient}s for the ADK. The default clients are cached per name so the
33+
* dispatcher and connection pool are reused across the ADK; a caller that supplies its own executor
34+
* gets a fresh, non-cached client it owns.
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 the shared {@link OkHttpClient} cached by {@code name}, using OkHttp's default
44+
* threading.
45+
*/
46+
public static OkHttpClient getOrCreateSharedHttpClient(String name) {
47+
return sharedClients.computeIfAbsent(name, unused -> new OkHttpClient());
48+
}
49+
50+
/**
51+
* Returns a new {@link OkHttpClient} whose dispatcher runs on {@code executorService}. Pass
52+
* {@link #daemonExecutor} so a standalone JVM can exit once work is done, or a container-managed
53+
* executor in a managed environment. The client is not cached: the caller owns the executor and
54+
* the returned client.
55+
*
56+
* @param executorService executor for the dispatcher.
57+
*/
58+
public static OkHttpClient createHttpClient(ExecutorService executorService) {
59+
return new OkHttpClient.Builder().dispatcher(new Dispatcher(executorService)).build();
60+
}
61+
62+
/**
63+
* Returns an unbounded pool of daemon threads, matching OkHttp's own dispatcher pool but with
64+
* daemon threads so a standalone JVM can exit once work is done. Managed container environments
65+
* should inject their own executor instead of calling this.
66+
*
67+
* @param name prefix for the dispatcher thread names.
68+
*/
69+
public static ExecutorService daemonExecutor(String name) {
70+
return new ThreadPoolExecutor(
71+
0, Integer.MAX_VALUE, 60L, SECONDS, new SynchronousQueue<>(), daemonThreadFactory(name));
72+
}
73+
74+
private static ThreadFactory daemonThreadFactory(String name) {
75+
AtomicInteger count = new AtomicInteger();
76+
return runnable -> {
77+
Thread thread = new Thread(runnable, name + "-" + count.incrementAndGet());
78+
thread.setDaemon(true);
79+
return thread;
80+
};
81+
}
82+
}

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

Lines changed: 73 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,52 @@ 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+
OkHttpClient client =
67+
executorService == null
68+
? HttpClientFactory.getOrCreateSharedHttpClient("GeminiApiClient")
69+
: HttpClientFactory.createHttpClient(executorService);
70+
return client
71+
.newBuilder()
72+
.connectTimeout(Duration.ZERO)
73+
.readTimeout(Duration.ZERO)
74+
.writeTimeout(Duration.ZERO)
75+
.build();
76+
}
77+
78+
private static Client buildApiKeyClient(
79+
String apiKey, @Nullable ExecutorService executorService) {
80+
return Client.builder()
81+
.apiKey(apiKey)
82+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
83+
.clientOptions(
84+
ClientOptions.builder().customHttpClient(getSharedHttpClient(executorService)).build())
85+
.build();
86+
}
87+
88+
private static Client buildVertexClient(
89+
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
90+
Client.Builder apiClientBuilder =
91+
Client.builder()
92+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
93+
.clientOptions(
94+
ClientOptions.builder()
95+
.customHttpClient(getSharedHttpClient(executorService))
96+
.build());
97+
vertexCredentials.project().ifPresent(apiClientBuilder::project);
98+
vertexCredentials.location().ifPresent(apiClientBuilder::location);
99+
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
100+
return apiClientBuilder.build();
101+
}
102+
103+
private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
104+
return Client.builder()
105+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
106+
.clientOptions(
107+
ClientOptions.builder().customHttpClient(getSharedHttpClient(executorService)).build())
108+
.build();
109+
}
110+
60111
static {
61112
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
62113
String languageLabel = "gl-java/" + JAVA_VERSION.value();
@@ -90,11 +141,7 @@ public Gemini(String modelName, Client apiClient) {
90141
public Gemini(String modelName, String apiKey) {
91142
super(modelName);
92143
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();
144+
this.apiClient = buildApiKeyClient(apiKey, null);
98145
}
99146

100147
/**
@@ -106,12 +153,7 @@ public Gemini(String modelName, String apiKey) {
106153
public Gemini(String modelName, VertexCredentials vertexCredentials) {
107154
super(modelName);
108155
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();
156+
this.apiClient = buildVertexClient(vertexCredentials, null);
115157
}
116158

117159
/**
@@ -131,6 +173,7 @@ public static class Builder {
131173
private Client apiClient;
132174
private String apiKey;
133175
private VertexCredentials vertexCredentials;
176+
private ExecutorService httpExecutorService;
134177

135178
private Builder() {}
136179

@@ -186,6 +229,22 @@ public Builder vertexCredentials(VertexCredentials vertexCredentials) {
186229
return this;
187230
}
188231

232+
/**
233+
* Sets the executor for the shared HTTP client's dispatcher. Pass {@link
234+
* HttpClientFactory#daemonExecutor} so a standalone or CLI JVM can exit once work is done, or a
235+
* container-managed executor in a managed environment. Applies only when the client is built
236+
* from an API key, Vertex credentials, or the default; it is ignored when an explicit {@link
237+
* #apiClient} is supplied.
238+
*
239+
* @param httpExecutorService The executor for HTTP dispatcher threads.
240+
* @return This builder.
241+
*/
242+
@CanIgnoreReturnValue
243+
public Builder httpExecutorService(ExecutorService httpExecutorService) {
244+
this.httpExecutorService = httpExecutorService;
245+
return this;
246+
}
247+
189248
/**
190249
* Builds the {@link Gemini} instance.
191250
*
@@ -198,15 +257,11 @@ public Gemini build() {
198257
if (apiClient != null) {
199258
return new Gemini(modelName, apiClient);
200259
} else if (apiKey != null) {
201-
return new Gemini(modelName, apiKey);
260+
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService));
202261
} else if (vertexCredentials != null) {
203-
return new Gemini(modelName, vertexCredentials);
262+
return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService));
204263
} else {
205-
return new Gemini(
206-
modelName,
207-
Client.builder()
208-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
209-
.build());
264+
return new Gemini(modelName, buildDefaultClient(httpExecutorService));
210265
}
211266
}
212267
}

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

Lines changed: 29 additions & 8 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,16 @@ 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
74+
* Returns the OkHttpClient whose connection pool and thread dispatcher back {@link
75+
* ChatCompletionsHttpClient} instances. Without an executor this is the shared client cached by
76+
* name; with one it is a fresh client the caller owns. Each instance forks it via {@link
7377
* OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking pools.
7478
*/
75-
private static final OkHttpClient SHARED_POOL_CLIENT = new OkHttpClient();
79+
private static OkHttpClient getSharedPoolClient(@Nullable ExecutorService executorService) {
80+
return executorService == null
81+
? HttpClientFactory.getOrCreateSharedHttpClient("ChatCompletionsHttpClient")
82+
: HttpClientFactory.createHttpClient(executorService);
83+
}
7684

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

122142
private ChatCompletionsHttpClient(HttpOptions httpOptions, OkHttpClient client) {
@@ -153,12 +173,13 @@ static ChatCompletionsHttpClient forTesting(HttpOptions httpOptions, OkHttpClien
153173
}
154174

155175
/**
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.
176+
* Builds the production OkHttpClient by forking the shared pool client so the connection pool and
177+
* dispatcher are reused across instances while applying per-instance timeouts.
158178
*/
159-
private static OkHttpClient buildClient(HttpOptions httpOptions) {
179+
private static OkHttpClient buildClient(
180+
HttpOptions httpOptions, @Nullable ExecutorService executorService) {
160181
Objects.requireNonNull(httpOptions, "httpOptions cannot be null");
161-
OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder();
182+
OkHttpClient.Builder builder = getSharedPoolClient(executorService).newBuilder();
162183
builder.connectTimeout(Duration.ZERO);
163184
builder.readTimeout(Duration.ZERO);
164185
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.getOrCreateSharedHttpClient("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)