chore(client): refactor closing / shutdown

Author: stainless-app[bot]

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

@@ -126,6 +126,8 @@ class BraintrustOkHttpClient private constructor() {
          * The executor to use for running [AsyncStreamResponse.Handler] callbacks.
          *
          * Defaults to a dedicated cached thread pool.
+         *
+         * This class takes ownership of the executor and shuts it down, if possible, when closed.
          */
         fun streamHandlerExecutor(streamHandlerExecutor: Executor) = apply {
             clientOptions.streamHandlerExecutor(streamHandlerExecutor)

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

@@ -126,6 +126,8 @@ class BraintrustOkHttpClientAsync private constructor() {
          * The executor to use for running [AsyncStreamResponse.Handler] callbacks.
          *
          * Defaults to a dedicated cached thread pool.
+         *
+         * This class takes ownership of the executor and shuts it down, if possible, when closed.
          */
         fun streamHandlerExecutor(streamHandlerExecutor: Executor) = apply {
             clientOptions.streamHandlerExecutor(streamHandlerExecutor)

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

@@ -172,7 +172,7 @@ class BraintrustClientAsyncImpl(private val clientOptions: ClientOptions) : Brai
 
     override fun evals(): EvalServiceAsync = evals
 
-    override fun close() = clientOptions.httpClient.close()
+    override fun close() = clientOptions.close()
 
     class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) :
         BraintrustClientAsync.WithRawResponse {

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

@@ -160,7 +160,7 @@ class BraintrustClientImpl(private val clientOptions: ClientOptions) : Braintrus
 
     override fun evals(): EvalService = evals
 
-    override fun close() = clientOptions.httpClient.close()
+    override fun close() = clientOptions.close()
 
     class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) :
         BraintrustClient.WithRawResponse {

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

@@ -13,6 +13,7 @@ import java.time.Clock
 import java.time.Duration
 import java.util.Optional
 import java.util.concurrent.Executor
+import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
 import java.util.concurrent.ThreadFactory
 import java.util.concurrent.atomic.AtomicLong
@@ -26,6 +27,8 @@ private constructor(
      * The HTTP client to use in the SDK.
      *
      * Use the one published in `braintrust-java-client-okhttp` or implement your own.
+     *
+     * This class takes ownership of the client and closes it when closed.
      */
     @get:JvmName("httpClient") val httpClient: HttpClient,
     /**
@@ -47,6 +50,8 @@ private constructor(
      * The executor to use for running [AsyncStreamResponse.Handler] callbacks.
      *
      * Defaults to a dedicated cached thread pool.
+     *
+     * This class takes ownership of the executor and shuts it down, if possible, when closed.
      */
     @get:JvmName("streamHandlerExecutor") val streamHandlerExecutor: Executor,
     /**
@@ -170,6 +175,8 @@ private constructor(
          * The HTTP client to use in the SDK.
          *
          * Use the one published in `braintrust-java-client-okhttp` or implement your own.
+         *
+         * This class takes ownership of the client and closes it when closed.
          */
         fun httpClient(httpClient: HttpClient) = apply {
             this.httpClient = PhantomReachableClosingHttpClient(httpClient)
@@ -198,9 +205,14 @@ private constructor(
          * The executor to use for running [AsyncStreamResponse.Handler] callbacks.
          *
          * Defaults to a dedicated cached thread pool.
+         *
+         * This class takes ownership of the executor and shuts it down, if possible, when closed.
          */
         fun streamHandlerExecutor(streamHandlerExecutor: Executor) = apply {
-            this.streamHandlerExecutor = streamHandlerExecutor
+            this.streamHandlerExecutor =
+                if (streamHandlerExecutor is ExecutorService)
+                    PhantomReachableExecutorService(streamHandlerExecutor)
+                else streamHandlerExecutor
         }
 
         /**
@@ -440,4 +452,19 @@ private constructor(
             )
         }
     }
+
+    /**
+     * Closes these client options, relinquishing any underlying resources.
+     *
+     * This is purposefully not inherited from [AutoCloseable] because the client options are
+     * long-lived and usually should not be synchronously closed via try-with-resources.
+     *
+     * It's also usually not necessary to call this method at all. the default client automatically
+     * releases threads and connections if they remain idle, but if you are writing an application
+     * that needs to aggressively release unused resources, then you may call this method.
+     */
+    fun close() {
+        httpClient.close()
+        (streamHandlerExecutor as? ExecutorService)?.shutdown()
+    }
 }

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

@@ -0,0 +1,58 @@
+package com.braintrustdata.api.core
+
+import java.util.concurrent.Callable
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Future
+import java.util.concurrent.TimeUnit
+
+/**
+ * A delegating wrapper around an [ExecutorService] that shuts it down once it's only phantom
+ * reachable.
+ *
+ * This class ensures the [ExecutorService] is shut down even if the user forgets to do it.
+ */
+internal class PhantomReachableExecutorService(private val executorService: ExecutorService) :
+    ExecutorService {
+    init {
+        closeWhenPhantomReachable(this) { executorService.shutdown() }
+    }
+
+    override fun execute(command: Runnable) = executorService.execute(command)
+
+    override fun shutdown() = executorService.shutdown()
+
+    override fun shutdownNow(): MutableList<Runnable> = executorService.shutdownNow()
+
+    override fun isShutdown(): Boolean = executorService.isShutdown
+
+    override fun isTerminated(): Boolean = executorService.isTerminated
+
+    override fun awaitTermination(timeout: Long, unit: TimeUnit): Boolean =
+        executorService.awaitTermination(timeout, unit)
+
+    override fun <T : Any?> submit(task: Callable<T>): Future<T> = executorService.submit(task)
+
+    override fun <T : Any?> submit(task: Runnable, result: T): Future<T> =
+        executorService.submit(task, result)
+
+    override fun submit(task: Runnable): Future<*> = executorService.submit(task)
+
+    override fun <T : Any?> invokeAll(
+        tasks: MutableCollection<out Callable<T>>
+    ): MutableList<Future<T>> = executorService.invokeAll(tasks)
+
+    override fun <T : Any?> invokeAll(
+        tasks: MutableCollection<out Callable<T>>,
+        timeout: Long,
+        unit: TimeUnit,
+    ): MutableList<Future<T>> = executorService.invokeAll(tasks, timeout, unit)
+
+    override fun <T : Any?> invokeAny(tasks: MutableCollection<out Callable<T>>): T =
+        executorService.invokeAny(tasks)
+
+    override fun <T : Any?> invokeAny(
+        tasks: MutableCollection<out Callable<T>>,
+        timeout: Long,
+        unit: TimeUnit,
+    ): T = executorService.invokeAny(tasks, timeout, unit)
+}