Skip to content

Commit 22a3be8

Browse files
feat(client): expose sleeper option
fix(client): ensure single timer is created per client
1 parent 6932929 commit 22a3be8

File tree

8 files changed

+158
-49
lines changed

8 files changed

+158
-49
lines changed

openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.openai.azure.AzureUrlPathMode
88
import com.openai.client.OpenAIClient
99
import com.openai.client.OpenAIClientImpl
1010
import com.openai.core.ClientOptions
11+
import com.openai.core.Sleeper
1112
import com.openai.core.Timeout
1213
import com.openai.core.http.AsyncStreamResponse
1314
import com.openai.core.http.Headers
@@ -135,6 +136,17 @@ class OpenAIOkHttpClient private constructor() {
135136
clientOptions.streamHandlerExecutor(streamHandlerExecutor)
136137
}
137138

139+
/**
140+
* The interface to use for delaying execution, like during retries.
141+
*
142+
* This is primarily useful for using fake delays in tests.
143+
*
144+
* Defaults to real execution delays.
145+
*
146+
* This class takes ownership of the sleeper and closes it when closed.
147+
*/
148+
fun sleeper(sleeper: Sleeper) = apply { clientOptions.sleeper(sleeper) }
149+
138150
/**
139151
* The clock to use for operations that require timing, like retries.
140152
*

openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.openai.azure.AzureUrlPathMode
88
import com.openai.client.OpenAIClientAsync
99
import com.openai.client.OpenAIClientAsyncImpl
1010
import com.openai.core.ClientOptions
11+
import com.openai.core.Sleeper
1112
import com.openai.core.Timeout
1213
import com.openai.core.http.AsyncStreamResponse
1314
import com.openai.core.http.Headers
@@ -135,6 +136,17 @@ class OpenAIOkHttpClientAsync private constructor() {
135136
clientOptions.streamHandlerExecutor(streamHandlerExecutor)
136137
}
137138

139+
/**
140+
* The interface to use for delaying execution, like during retries.
141+
*
142+
* This is primarily useful for using fake delays in tests.
143+
*
144+
* Defaults to real execution delays.
145+
*
146+
* This class takes ownership of the sleeper and closes it when closed.
147+
*/
148+
fun sleeper(sleeper: Sleeper) = apply { clientOptions.sleeper(sleeper) }
149+
138150
/**
139151
* The clock to use for operations that require timing, like retries.
140152
*

openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,16 @@ private constructor(
6060
* This class takes ownership of the executor and shuts it down, if possible, when closed.
6161
*/
6262
@get:JvmName("streamHandlerExecutor") val streamHandlerExecutor: Executor,
63+
/**
64+
* The interface to use for delaying execution, like during retries.
65+
*
66+
* This is primarily useful for using fake delays in tests.
67+
*
68+
* Defaults to real execution delays.
69+
*
70+
* This class takes ownership of the sleeper and closes it when closed.
71+
*/
72+
@get:JvmName("sleeper") val sleeper: Sleeper,
6373
/**
6474
* The clock to use for operations that require timing, like retries.
6575
*
@@ -162,6 +172,7 @@ private constructor(
162172
private var checkJacksonVersionCompatibility: Boolean = true
163173
private var jsonMapper: JsonMapper = jsonMapper()
164174
private var streamHandlerExecutor: Executor? = null
175+
private var sleeper: Sleeper? = null
165176
private var clock: Clock = Clock.systemUTC()
166177
private var baseUrl: String? = null
167178
private var headers: Headers.Builder = Headers.builder()
@@ -182,6 +193,7 @@ private constructor(
182193
checkJacksonVersionCompatibility = clientOptions.checkJacksonVersionCompatibility
183194
jsonMapper = clientOptions.jsonMapper
184195
streamHandlerExecutor = clientOptions.streamHandlerExecutor
196+
sleeper = clientOptions.sleeper
185197
clock = clientOptions.clock
186198
baseUrl = clientOptions.baseUrl
187199
headers = clientOptions.headers.toBuilder()
@@ -241,6 +253,17 @@ private constructor(
241253
else streamHandlerExecutor
242254
}
243255

256+
/**
257+
* The interface to use for delaying execution, like during retries.
258+
*
259+
* This is primarily useful for using fake delays in tests.
260+
*
261+
* Defaults to real execution delays.
262+
*
263+
* This class takes ownership of the sleeper and closes it when closed.
264+
*/
265+
fun sleeper(sleeper: Sleeper) = apply { this.sleeper = PhantomReachableSleeper(sleeper) }
266+
244267
/**
245268
* The clock to use for operations that require timing, like retries.
246269
*
@@ -479,6 +502,25 @@ private constructor(
479502
*/
480503
fun build(): ClientOptions {
481504
val httpClient = checkRequired("httpClient", httpClient)
505+
val streamHandlerExecutor =
506+
streamHandlerExecutor
507+
?: PhantomReachableExecutorService(
508+
Executors.newCachedThreadPool(
509+
object : ThreadFactory {
510+
511+
private val threadFactory: ThreadFactory =
512+
Executors.defaultThreadFactory()
513+
private val count = AtomicLong(0)
514+
515+
override fun newThread(runnable: Runnable): Thread =
516+
threadFactory.newThread(runnable).also {
517+
it.name =
518+
"openai-stream-handler-thread-${count.getAndIncrement()}"
519+
}
520+
}
521+
)
522+
)
523+
val sleeper = sleeper ?: PhantomReachableSleeper(DefaultSleeper())
482524
val credential = checkRequired("credential", credential)
483525

484526
val headers = Headers.builder()
@@ -530,26 +572,14 @@ private constructor(
530572
httpClient,
531573
RetryingHttpClient.builder()
532574
.httpClient(httpClient)
575+
.sleeper(sleeper)
533576
.clock(clock)
534577
.maxRetries(maxRetries)
535578
.build(),
536579
checkJacksonVersionCompatibility,
537580
jsonMapper,
538-
streamHandlerExecutor
539-
?: Executors.newCachedThreadPool(
540-
object : ThreadFactory {
541-
542-
private val threadFactory: ThreadFactory =
543-
Executors.defaultThreadFactory()
544-
private val count = AtomicLong(0)
545-
546-
override fun newThread(runnable: Runnable): Thread =
547-
threadFactory.newThread(runnable).also {
548-
it.name =
549-
"openai-stream-handler-thread-${count.getAndIncrement()}"
550-
}
551-
}
552-
),
581+
streamHandlerExecutor,
582+
sleeper,
553583
clock,
554584
baseUrl,
555585
headers.build(),
@@ -580,5 +610,6 @@ private constructor(
580610
fun close() {
581611
httpClient.close()
582612
(streamHandlerExecutor as? ExecutorService)?.shutdown()
613+
sleeper.close()
583614
}
584615
}
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package com.openai.core
2+
3+
import java.time.Duration
4+
import java.util.Timer
5+
import java.util.TimerTask
6+
import java.util.concurrent.CompletableFuture
7+
8+
class DefaultSleeper : Sleeper {
9+
10+
private val timer = Timer("DefaultSleeper", true)
11+
12+
override fun sleep(duration: Duration) = Thread.sleep(duration.toMillis())
13+
14+
override fun sleepAsync(duration: Duration): CompletableFuture<Void> {
15+
val future = CompletableFuture<Void>()
16+
timer.schedule(
17+
object : TimerTask() {
18+
override fun run() {
19+
future.complete(null)
20+
}
21+
},
22+
duration.toMillis(),
23+
)
24+
return future
25+
}
26+
27+
override fun close() = timer.cancel()
28+
}
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package com.openai.core
2+
3+
import java.time.Duration
4+
import java.util.concurrent.CompletableFuture
5+
6+
/**
7+
* A delegating wrapper around a [Sleeper] that closes it once it's only phantom reachable.
8+
*
9+
* This class ensures the [Sleeper] is closed even if the user forgets to do it.
10+
*/
11+
internal class PhantomReachableSleeper(private val sleeper: Sleeper) : Sleeper {
12+
13+
init {
14+
closeWhenPhantomReachable(this, sleeper)
15+
}
16+
17+
override fun sleep(duration: Duration) = sleeper.sleep(duration)
18+
19+
override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
20+
sleeper.sleepAsync(duration)
21+
22+
override fun close() = sleeper.close()
23+
}
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package com.openai.core
2+
3+
import java.time.Duration
4+
import java.util.concurrent.CompletableFuture
5+
6+
/**
7+
* An interface for delaying execution for a specified amount of time.
8+
*
9+
* Useful for testing and cleaning up resources.
10+
*/
11+
interface Sleeper : AutoCloseable {
12+
13+
/** Synchronously pauses execution for the given [duration]. */
14+
fun sleep(duration: Duration)
15+
16+
/** Asynchronously pauses execution for the given [duration]. */
17+
fun sleepAsync(duration: Duration): CompletableFuture<Void>
18+
19+
/** Overridden from [AutoCloseable] to not have a checked exception in its signature. */
20+
override fun close()
21+
}

openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt

Lines changed: 9 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package com.openai.core.http
22

3+
import com.openai.core.DefaultSleeper
34
import com.openai.core.RequestOptions
5+
import com.openai.core.Sleeper
46
import com.openai.core.checkRequired
57
import com.openai.errors.OpenAIIoException
68
import com.openai.errors.OpenAIRetryableException
@@ -11,8 +13,6 @@ import java.time.OffsetDateTime
1113
import java.time.format.DateTimeFormatter
1214
import java.time.format.DateTimeParseException
1315
import java.time.temporal.ChronoUnit
14-
import java.util.Timer
15-
import java.util.TimerTask
1616
import java.util.UUID
1717
import java.util.concurrent.CompletableFuture
1818
import java.util.concurrent.ThreadLocalRandom
@@ -130,7 +130,10 @@ private constructor(
130130
return executeWithRetries(modifiedRequest, requestOptions)
131131
}
132132

133-
override fun close() = httpClient.close()
133+
override fun close() {
134+
httpClient.close()
135+
sleeper.close()
136+
}
134137

135138
private fun isRetryable(request: HttpRequest): Boolean =
136139
// Some requests, such as when a request body is being streamed, cannot be retried because
@@ -235,33 +238,14 @@ private constructor(
235238
class Builder internal constructor() {
236239

237240
private var httpClient: HttpClient? = null
238-
private var sleeper: Sleeper =
239-
object : Sleeper {
240-
241-
private val timer = Timer("RetryingHttpClient", true)
242-
243-
override fun sleep(duration: Duration) = Thread.sleep(duration.toMillis())
244-
245-
override fun sleepAsync(duration: Duration): CompletableFuture<Void> {
246-
val future = CompletableFuture<Void>()
247-
timer.schedule(
248-
object : TimerTask() {
249-
override fun run() {
250-
future.complete(null)
251-
}
252-
},
253-
duration.toMillis(),
254-
)
255-
return future
256-
}
257-
}
241+
private var sleeper: Sleeper? = null
258242
private var clock: Clock = Clock.systemUTC()
259243
private var maxRetries: Int = 2
260244
private var idempotencyHeader: String? = null
261245

262246
fun httpClient(httpClient: HttpClient) = apply { this.httpClient = httpClient }
263247

264-
@JvmSynthetic internal fun sleeper(sleeper: Sleeper) = apply { this.sleeper = sleeper }
248+
fun sleeper(sleeper: Sleeper) = apply { this.sleeper = sleeper }
265249

266250
fun clock(clock: Clock) = apply { this.clock = clock }
267251

@@ -272,17 +256,10 @@ private constructor(
272256
fun build(): HttpClient =
273257
RetryingHttpClient(
274258
checkRequired("httpClient", httpClient),
275-
sleeper,
259+
sleeper ?: DefaultSleeper(),
276260
clock,
277261
maxRetries,
278262
idempotencyHeader,
279263
)
280264
}
281-
282-
internal interface Sleeper {
283-
284-
fun sleep(duration: Duration)
285-
286-
fun sleepAsync(duration: Duration): CompletableFuture<Void>
287-
}
288265
}

openai-java-core/src/test/kotlin/com/openai/core/http/RetryingHttpClientTest.kt

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import com.github.tomakehurst.wiremock.junit5.WireMockTest
66
import com.github.tomakehurst.wiremock.stubbing.Scenario
77
import com.openai.client.okhttp.OkHttpClient
88
import com.openai.core.RequestOptions
9+
import com.openai.core.Sleeper
910
import com.openai.errors.OpenAIRetryableException
1011
import java.io.InputStream
1112
import java.time.Duration
@@ -294,12 +295,14 @@ internal class RetryingHttpClientTest {
294295
.httpClient(failingHttpClient)
295296
.maxRetries(2)
296297
.sleeper(
297-
object : RetryingHttpClient.Sleeper {
298+
object : Sleeper {
298299

299300
override fun sleep(duration: Duration) {}
300301

301302
override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
302303
CompletableFuture.completedFuture(null)
304+
305+
override fun close() {}
303306
}
304307
)
305308
.build()
@@ -333,12 +336,14 @@ internal class RetryingHttpClientTest {
333336
.httpClient(httpClient)
334337
// Use a no-op `Sleeper` to make the test fast.
335338
.sleeper(
336-
object : RetryingHttpClient.Sleeper {
339+
object : Sleeper {
337340

338341
override fun sleep(duration: Duration) {}
339342

340343
override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
341344
CompletableFuture.completedFuture(null)
345+
346+
override fun close() {}
342347
}
343348
)
344349

0 commit comments

Comments
 (0)