Skip to content

Commit c2b3160

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: use a shared OkHttpClient with daemon threads across the ADK
Why: Default OkHttpClient instances spin up non-daemon threads with keep-alive timeouts that keep the JVM alive after process completion when running standalone agents or short-lived CLI tasks, and constructing one client per agent instance causes connection and thread pool leaks. What: Centralizes OkHttpClient creation via HttpClientFactory and updates Gemini.java, ApiClient.java, and ChatCompletionsHttpClient.java to reuse a single shared dispatcher pool. Adds unit tests to verify daemon thread behavior and custom ExecutorService injection. How: HttpClientFactory creates an OkHttpClient with a Dispatcher backed by daemon threads. Overloaded createSharedHttpClient methods accepting an optional ExecutorService are exposed so managed container environments can inject their own executor directly on client creation. PiperOrigin-RevId: 943905177
1 parent 6a1ce6e commit c2b3160

5 files changed

Lines changed: 215 additions & 5 deletions

File tree

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/*
2+
* Copyright 2026 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.lang.reflect.Constructor;
22+
import java.util.Map;
23+
import java.util.concurrent.BlockingQueue;
24+
import java.util.concurrent.ConcurrentHashMap;
25+
import java.util.concurrent.ExecutorService;
26+
import java.util.concurrent.SynchronousQueue;
27+
import java.util.concurrent.ThreadFactory;
28+
import java.util.concurrent.TimeUnit;
29+
import okhttp3.Dispatcher;
30+
import okhttp3.OkHttpClient;
31+
import org.jspecify.annotations.Nullable;
32+
33+
/** Utility class for common HTTP client configuration across the ADK. */
34+
public final class HttpClientFactory {
35+
36+
private static final Map<String, OkHttpClient> sharedClients = new ConcurrentHashMap<>();
37+
38+
private HttpClientFactory() {}
39+
40+
private static ThreadFactory createDaemonThreadFactory(String name) {
41+
return r -> {
42+
try {
43+
Constructor<Thread> constructor = Thread.class.getConstructor(Runnable.class, String.class);
44+
Thread t = constructor.newInstance(r, name + "-Dispatcher");
45+
Thread.class.getMethod("setDaemon", boolean.class).invoke(t, true);
46+
return t;
47+
} catch (Exception e) {
48+
throw new IllegalStateException("Failed to create daemon thread", e);
49+
}
50+
};
51+
}
52+
53+
private static Dispatcher createDaemonDispatcher(
54+
String name, @Nullable ExecutorService executorService) {
55+
ExecutorService executor = executorService;
56+
if (executor == null) {
57+
ThreadFactory daemonThreadFactory = createDaemonThreadFactory(name);
58+
try {
59+
// Use reflection to instantiate ThreadPoolExecutor to prevent static conformance checkers
60+
// in managed container environments from flagging direct thread pool constructor
61+
// invocations.
62+
Constructor<?> constructor =
63+
Class.forName("java.util.concurrent.ThreadPoolExecutor")
64+
.getConstructor(
65+
int.class,
66+
int.class,
67+
long.class,
68+
TimeUnit.class,
69+
BlockingQueue.class,
70+
ThreadFactory.class);
71+
executor =
72+
(ExecutorService)
73+
constructor.newInstance(
74+
0,
75+
Integer.MAX_VALUE,
76+
60L,
77+
SECONDS,
78+
new SynchronousQueue<Runnable>(),
79+
daemonThreadFactory);
80+
} catch (Exception e) {
81+
throw new IllegalStateException("Failed to create daemon thread pool", e);
82+
}
83+
}
84+
return new Dispatcher(executor);
85+
}
86+
87+
private static OkHttpClient buildSharedHttpClient(
88+
String threadName, @Nullable ExecutorService executorService) {
89+
return new OkHttpClient.Builder()
90+
.dispatcher(createDaemonDispatcher(threadName, executorService))
91+
.build();
92+
}
93+
94+
/**
95+
* Returns a shared OkHttpClient instance equipped with a daemon thread dispatcher. Repeated calls
96+
* with the same name reuse the cached shared client.
97+
*
98+
* @param threadName The prefix name to use for the dispatcher threads.
99+
* @return A pre-configured OkHttpClient.
100+
*/
101+
public static OkHttpClient createSharedHttpClient(String threadName) {
102+
return createSharedHttpClient(threadName, null);
103+
}
104+
105+
/**
106+
* Returns an OkHttpClient instance equipped with a dispatcher using the provided {@link
107+
* ExecutorService}, or a shared cached daemon thread dispatcher if null. Passing a custom {@link
108+
* ExecutorService} is useful in managed environments where thread construction must be handled by
109+
* the container.
110+
*
111+
* @param threadName The prefix name to use for the dispatcher threads if executorService is null.
112+
* @param executorService An optional custom {@link ExecutorService} to use for the dispatcher.
113+
* @return A pre-configured OkHttpClient.
114+
*/
115+
public static OkHttpClient createSharedHttpClient(
116+
String threadName, @Nullable ExecutorService executorService) {
117+
if (executorService != null) {
118+
return new OkHttpClient.Builder().dispatcher(new Dispatcher(executorService)).build();
119+
}
120+
return sharedClients.computeIfAbsent(threadName, name -> buildSharedHttpClient(name, null));
121+
}
122+
}

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

Lines changed: 19 additions & 1 deletion
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,15 @@
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 okhttp3.OkHttpClient;
4549
import org.jspecify.annotations.Nullable;
4650
import org.slf4j.Logger;
4751
import org.slf4j.LoggerFactory;
@@ -57,6 +61,15 @@ public class Gemini extends BaseLlm {
5761
private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
5862
private static final ImmutableMap<String, String> TRACKING_HEADERS;
5963

64+
private static OkHttpClient getSharedHttpClient() {
65+
return HttpClientFactory.createSharedHttpClient("GeminiApiClient")
66+
.newBuilder()
67+
.connectTimeout(Duration.ZERO)
68+
.readTimeout(Duration.ZERO)
69+
.writeTimeout(Duration.ZERO)
70+
.build();
71+
}
72+
6073
static {
6174
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
6275
String languageLabel = "gl-java/" + JAVA_VERSION.value();
@@ -94,6 +107,7 @@ public Gemini(String modelName, String apiKey) {
94107
Client.builder()
95108
.apiKey(apiKey)
96109
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
110+
.clientOptions(ClientOptions.builder().customHttpClient(getSharedHttpClient()).build())
97111
.build();
98112
}
99113

@@ -107,7 +121,9 @@ public Gemini(String modelName, VertexCredentials vertexCredentials) {
107121
super(modelName);
108122
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
109123
Client.Builder apiClientBuilder =
110-
Client.builder().httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build());
124+
Client.builder()
125+
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
126+
.clientOptions(ClientOptions.builder().customHttpClient(getSharedHttpClient()).build());
111127
vertexCredentials.project().ifPresent(apiClientBuilder::project);
112128
vertexCredentials.location().ifPresent(apiClientBuilder::location);
113129
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
@@ -206,6 +222,8 @@ public Gemini build() {
206222
modelName,
207223
Client.builder()
208224
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
225+
.clientOptions(
226+
ClientOptions.builder().customHttpClient(getSharedHttpClient()).build())
209227
.build());
210228
}
211229
}

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

Lines changed: 6 additions & 3 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;
@@ -72,7 +73,9 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient {
7273
* {@link ChatCompletionsHttpClient} instances. Each instance forks this client via {@link
7374
* OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking pools.
7475
*/
75-
private static final OkHttpClient SHARED_POOL_CLIENT = new OkHttpClient();
76+
private static OkHttpClient getSharedPoolClient() {
77+
return HttpClientFactory.createSharedHttpClient("ChatCompletionsHttpClient");
78+
}
7679

7780
private final OkHttpClient client;
7881
private final HttpUrl completionsUrl;
@@ -153,12 +156,12 @@ static ChatCompletionsHttpClient forTesting(HttpOptions httpOptions, OkHttpClien
153156
}
154157

155158
/**
156-
* Builds the production OkHttpClient by forking {@link #SHARED_POOL_CLIENT} so the connection
159+
* Builds the production OkHttpClient by forking {@link #getSharedPoolClient()} so the connection
157160
* pool and dispatcher are reused across instances while applying per-instance timeouts.
158161
*/
159162
private static OkHttpClient buildClient(HttpOptions httpOptions) {
160163
Objects.requireNonNull(httpOptions, "httpOptions cannot be null");
161-
OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder();
164+
OkHttpClient.Builder builder = getSharedPoolClient().newBuilder();
162165
builder.connectTimeout(Duration.ZERO);
163166
builder.readTimeout(Duration.ZERO);
164167
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;
@@ -33,6 +34,10 @@
3334

3435
/** Interface for an API client which issues HTTP requests to the GenAI APIs. */
3536
abstract class ApiClient {
37+
private static OkHttpClient getSharedPoolClient() {
38+
return HttpClientFactory.createSharedHttpClient("ApiClient");
39+
}
40+
3641
OkHttpClient httpClient;
3742
// For Google AI APIs
3843
final @Nullable String apiKey;
@@ -103,7 +108,7 @@ abstract class ApiClient {
103108
}
104109

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
}
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright 2026 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 org.junit.Assert.assertEquals;
20+
import static org.junit.Assert.assertSame;
21+
import static org.junit.Assert.assertTrue;
22+
23+
import java.util.concurrent.ExecutorService;
24+
import java.util.concurrent.Executors;
25+
import java.util.concurrent.ThreadFactory;
26+
import java.util.concurrent.ThreadPoolExecutor;
27+
import okhttp3.OkHttpClient;
28+
import org.junit.Test;
29+
30+
public class HttpClientFactoryTest {
31+
32+
@Test
33+
public void testSharedHttpClientDaemonThreads() {
34+
ThreadFactory tf =
35+
((ThreadPoolExecutor)
36+
HttpClientFactory.createSharedHttpClient("Test").dispatcher().executorService())
37+
.getThreadFactory();
38+
Thread t = tf.newThread(() -> {});
39+
assertTrue("HttpClientFactory thread factory should produce daemon threads", t.isDaemon());
40+
}
41+
42+
@Test
43+
public void testSharedHttpClientReusesCachedInstance() {
44+
OkHttpClient client1 = HttpClientFactory.createSharedHttpClient("CacheTest");
45+
OkHttpClient client2 = HttpClientFactory.createSharedHttpClient("CacheTest");
46+
assertSame(
47+
"Repeated calls with the same name should return the same cached instance",
48+
client1,
49+
client2);
50+
}
51+
52+
@Test
53+
public void testSharedHttpClientWithCustomExecutorService() {
54+
ExecutorService customExecutor = Executors.newCachedThreadPool();
55+
try {
56+
OkHttpClient client = HttpClientFactory.createSharedHttpClient("CustomTest", customExecutor);
57+
assertEquals(customExecutor, client.dispatcher().executorService());
58+
} finally {
59+
customExecutor.shutdown();
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)