feat(client): expose sleeper option

fix(client): ensure single timer is created per client

Author: stainless-app[bot]

braintrust-java-client-okhttp/src/main/kotlin/com/braintrustdata/api/client/okhttp/BraintrustOkHttpClient.kt

@@ -5,6 +5,7 @@ package com.braintrustdata.api.client.okhttp
 import com.braintrustdata.api.client.BraintrustClient
 import com.braintrustdata.api.client.BraintrustClientImpl
 import com.braintrustdata.api.core.ClientOptions
+import com.braintrustdata.api.core.Sleeper
 import com.braintrustdata.api.core.Timeout
 import com.braintrustdata.api.core.http.AsyncStreamResponse
 import com.braintrustdata.api.core.http.Headers
@@ -133,6 +134,17 @@ class BraintrustOkHttpClient private constructor() {
             clientOptions.streamHandlerExecutor(streamHandlerExecutor)
         }
 
+        /**
+         * The interface to use for delaying execution, like during retries.
+         *
+         * This is primarily useful for using fake delays in tests.
+         *
+         * Defaults to real execution delays.
+         *
+         * This class takes ownership of the sleeper and closes it when closed.
+         */
+        fun sleeper(sleeper: Sleeper) = apply { clientOptions.sleeper(sleeper) }
+
         /**
          * The clock to use for operations that require timing, like retries.
          *

braintrust-java-client-okhttp/src/main/kotlin/com/braintrustdata/api/client/okhttp/BraintrustOkHttpClientAsync.kt

@@ -5,6 +5,7 @@ package com.braintrustdata.api.client.okhttp
 import com.braintrustdata.api.client.BraintrustClientAsync
 import com.braintrustdata.api.client.BraintrustClientAsyncImpl
 import com.braintrustdata.api.core.ClientOptions
+import com.braintrustdata.api.core.Sleeper
 import com.braintrustdata.api.core.Timeout
 import com.braintrustdata.api.core.http.AsyncStreamResponse
 import com.braintrustdata.api.core.http.Headers
@@ -133,6 +134,17 @@ class BraintrustOkHttpClientAsync private constructor() {
             clientOptions.streamHandlerExecutor(streamHandlerExecutor)
         }
 
+        /**
+         * The interface to use for delaying execution, like during retries.
+         *
+         * This is primarily useful for using fake delays in tests.
+         *
+         * Defaults to real execution delays.
+         *
+         * This class takes ownership of the sleeper and closes it when closed.
+         */
+        fun sleeper(sleeper: Sleeper) = apply { clientOptions.sleeper(sleeper) }
+
         /**
          * The clock to use for operations that require timing, like retries.
          *

braintrust-java-core/src/main/kotlin/com/braintrustdata/api/core/ClientOptions.kt

@@ -54,6 +54,16 @@ private constructor(
      * This class takes ownership of the executor and shuts it down, if possible, when closed.
      */
     @get:JvmName("streamHandlerExecutor") val streamHandlerExecutor: Executor,
+    /**
+     * The interface to use for delaying execution, like during retries.
+     *
+     * This is primarily useful for using fake delays in tests.
+     *
+     * Defaults to real execution delays.
+     *
+     * This class takes ownership of the sleeper and closes it when closed.
+     */
+    @get:JvmName("sleeper") val sleeper: Sleeper,
     /**
      * The clock to use for operations that require timing, like retries.
      *
@@ -146,6 +156,7 @@ private constructor(
         private var checkJacksonVersionCompatibility: Boolean = true
         private var jsonMapper: JsonMapper = jsonMapper()
         private var streamHandlerExecutor: Executor? = null
+        private var sleeper: Sleeper? = null
         private var clock: Clock = Clock.systemUTC()
         private var baseUrl: String? = null
         private var headers: Headers.Builder = Headers.builder()
@@ -161,6 +172,7 @@ private constructor(
             checkJacksonVersionCompatibility = clientOptions.checkJacksonVersionCompatibility
             jsonMapper = clientOptions.jsonMapper
             streamHandlerExecutor = clientOptions.streamHandlerExecutor
+            sleeper = clientOptions.sleeper
             clock = clientOptions.clock
             baseUrl = clientOptions.baseUrl
             headers = clientOptions.headers.toBuilder()
@@ -215,6 +227,17 @@ private constructor(
                 else streamHandlerExecutor
         }
 
+        /**
+         * The interface to use for delaying execution, like during retries.
+         *
+         * This is primarily useful for using fake delays in tests.
+         *
+         * Defaults to real execution delays.
+         *
+         * This class takes ownership of the sleeper and closes it when closed.
+         */
+        fun sleeper(sleeper: Sleeper) = apply { this.sleeper = PhantomReachableSleeper(sleeper) }
+
         /**
          * The clock to use for operations that require timing, like retries.
          *
@@ -399,6 +422,25 @@ private constructor(
          */
         fun build(): ClientOptions {
             val httpClient = checkRequired("httpClient", httpClient)
+            val streamHandlerExecutor =
+                streamHandlerExecutor
+                    ?: PhantomReachableExecutorService(
+                        Executors.newCachedThreadPool(
+                            object : ThreadFactory {
+
+                                private val threadFactory: ThreadFactory =
+                                    Executors.defaultThreadFactory()
+                                private val count = AtomicLong(0)
+
+                                override fun newThread(runnable: Runnable): Thread =
+                                    threadFactory.newThread(runnable).also {
+                                        it.name =
+                                            "braintrust-stream-handler-thread-${count.getAndIncrement()}"
+                                    }
+                            }
+                        )
+                    )
+            val sleeper = sleeper ?: PhantomReachableSleeper(DefaultSleeper())
 
             val headers = Headers.builder()
             val queryParams = QueryParams.builder()
@@ -421,26 +463,14 @@ private constructor(
                 httpClient,
                 RetryingHttpClient.builder()
                     .httpClient(httpClient)
+                    .sleeper(sleeper)
                     .clock(clock)
                     .maxRetries(maxRetries)
                     .build(),
                 checkJacksonVersionCompatibility,
                 jsonMapper,
-                streamHandlerExecutor
-                    ?: Executors.newCachedThreadPool(
-                        object : ThreadFactory {
-
-                            private val threadFactory: ThreadFactory =
-                                Executors.defaultThreadFactory()
-                            private val count = AtomicLong(0)
-
-                            override fun newThread(runnable: Runnable): Thread =
-                                threadFactory.newThread(runnable).also {
-                                    it.name =
-                                        "braintrust-stream-handler-thread-${count.getAndIncrement()}"
-                                }
-                        }
-                    ),
+                streamHandlerExecutor,
+                sleeper,
                 clock,
                 baseUrl,
                 headers.build(),
@@ -466,5 +496,6 @@ private constructor(
     fun close() {
         httpClient.close()
         (streamHandlerExecutor as? ExecutorService)?.shutdown()
+        sleeper.close()
     }
 }

braintrust-java-core/src/main/kotlin/com/braintrustdata/api/core/DefaultSleeper.kt

@@ -0,0 +1,28 @@
+package com.braintrustdata.api.core
+
+import java.time.Duration
+import java.util.Timer
+import java.util.TimerTask
+import java.util.concurrent.CompletableFuture
+
+class DefaultSleeper : Sleeper {
+
+    private val timer = Timer("DefaultSleeper", true)
+
+    override fun sleep(duration: Duration) = Thread.sleep(duration.toMillis())
+
+    override fun sleepAsync(duration: Duration): CompletableFuture<Void> {
+        val future = CompletableFuture<Void>()
+        timer.schedule(
+            object : TimerTask() {
+                override fun run() {
+                    future.complete(null)
+                }
+            },
+            duration.toMillis(),
+        )
+        return future
+    }
+
+    override fun close() = timer.cancel()
+}

braintrust-java-core/src/main/kotlin/com/braintrustdata/api/core/PhantomReachableSleeper.kt

@@ -0,0 +1,23 @@
+package com.braintrustdata.api.core
+
+import java.time.Duration
+import java.util.concurrent.CompletableFuture
+
+/**
+ * A delegating wrapper around a [Sleeper] that closes it once it's only phantom reachable.
+ *
+ * This class ensures the [Sleeper] is closed even if the user forgets to do it.
+ */
+internal class PhantomReachableSleeper(private val sleeper: Sleeper) : Sleeper {
+
+    init {
+        closeWhenPhantomReachable(this, sleeper)
+    }
+
+    override fun sleep(duration: Duration) = sleeper.sleep(duration)
+
+    override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
+        sleeper.sleepAsync(duration)
+
+    override fun close() = sleeper.close()
+}

braintrust-java-core/src/main/kotlin/com/braintrustdata/api/core/Sleeper.kt

@@ -0,0 +1,21 @@
+package com.braintrustdata.api.core
+
+import java.time.Duration
+import java.util.concurrent.CompletableFuture
+
+/**
+ * An interface for delaying execution for a specified amount of time.
+ *
+ * Useful for testing and cleaning up resources.
+ */
+interface Sleeper : AutoCloseable {
+
+    /** Synchronously pauses execution for the given [duration]. */
+    fun sleep(duration: Duration)
+
+    /** Asynchronously pauses execution for the given [duration]. */
+    fun sleepAsync(duration: Duration): CompletableFuture<Void>
+
+    /** Overridden from [AutoCloseable] to not have a checked exception in its signature. */
+    override fun close()
+}

braintrust-java-core/src/main/kotlin/com/braintrustdata/api/core/http/RetryingHttpClient.kt

@@ -1,6 +1,8 @@
 package com.braintrustdata.api.core.http
 
+import com.braintrustdata.api.core.DefaultSleeper
 import com.braintrustdata.api.core.RequestOptions
+import com.braintrustdata.api.core.Sleeper
 import com.braintrustdata.api.core.checkRequired
 import com.braintrustdata.api.errors.BraintrustIoException
 import com.braintrustdata.api.errors.BraintrustRetryableException
@@ -11,8 +13,6 @@ import java.time.OffsetDateTime
 import java.time.format.DateTimeFormatter
 import java.time.format.DateTimeParseException
 import java.time.temporal.ChronoUnit
-import java.util.Timer
-import java.util.TimerTask
 import java.util.UUID
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ThreadLocalRandom
@@ -130,7 +130,10 @@ private constructor(
         return executeWithRetries(modifiedRequest, requestOptions)
     }
 
-    override fun close() = httpClient.close()
+    override fun close() {
+        httpClient.close()
+        sleeper.close()
+    }
 
     private fun isRetryable(request: HttpRequest): Boolean =
         // Some requests, such as when a request body is being streamed, cannot be retried because
@@ -235,33 +238,14 @@ private constructor(
     class Builder internal constructor() {
 
         private var httpClient: HttpClient? = null
-        private var sleeper: Sleeper =
-            object : Sleeper {
-
-                private val timer = Timer("RetryingHttpClient", true)
-
-                override fun sleep(duration: Duration) = Thread.sleep(duration.toMillis())
-
-                override fun sleepAsync(duration: Duration): CompletableFuture<Void> {
-                    val future = CompletableFuture<Void>()
-                    timer.schedule(
-                        object : TimerTask() {
-                            override fun run() {
-                                future.complete(null)
-                            }
-                        },
-                        duration.toMillis(),
-                    )
-                    return future
-                }
-            }
+        private var sleeper: Sleeper? = null
         private var clock: Clock = Clock.systemUTC()
         private var maxRetries: Int = 2
         private var idempotencyHeader: String? = null
 
         fun httpClient(httpClient: HttpClient) = apply { this.httpClient = httpClient }
 
-        @JvmSynthetic internal fun sleeper(sleeper: Sleeper) = apply { this.sleeper = sleeper }
+        fun sleeper(sleeper: Sleeper) = apply { this.sleeper = sleeper }
 
         fun clock(clock: Clock) = apply { this.clock = clock }
 
@@ -272,17 +256,10 @@ private constructor(
         fun build(): HttpClient =
             RetryingHttpClient(
                 checkRequired("httpClient", httpClient),
-                sleeper,
+                sleeper ?: DefaultSleeper(),
                 clock,
                 maxRetries,
                 idempotencyHeader,
             )
     }
-
-    internal interface Sleeper {
-
-        fun sleep(duration: Duration)
-
-        fun sleepAsync(duration: Duration): CompletableFuture<Void>
-    }
 }

braintrust-java-core/src/test/kotlin/com/braintrustdata/api/core/http/RetryingHttpClientTest.kt

@@ -2,6 +2,7 @@ package com.braintrustdata.api.core.http
 
 import com.braintrustdata.api.client.okhttp.OkHttpClient
 import com.braintrustdata.api.core.RequestOptions
+import com.braintrustdata.api.core.Sleeper
 import com.braintrustdata.api.errors.BraintrustRetryableException
 import com.github.tomakehurst.wiremock.client.WireMock.*
 import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo
@@ -294,12 +295,14 @@ internal class RetryingHttpClientTest {
                 .httpClient(failingHttpClient)
                 .maxRetries(2)
                 .sleeper(
-                    object : RetryingHttpClient.Sleeper {
+                    object : Sleeper {
 
                         override fun sleep(duration: Duration) {}
 
                         override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
                             CompletableFuture.completedFuture(null)
+
+                        override fun close() {}
                     }
                 )
                 .build()
@@ -333,12 +336,14 @@ internal class RetryingHttpClientTest {
             .httpClient(httpClient)
             // Use a no-op `Sleeper` to make the test fast.
             .sleeper(
-                object : RetryingHttpClient.Sleeper {
+                object : Sleeper {
 
                     override fun sleep(duration: Duration) {}
 
                     override fun sleepAsync(duration: Duration): CompletableFuture<Void> =
                         CompletableFuture.completedFuture(null)
+
+                    override fun close() {}
                 }
             )