Propagate async cancellation through retries

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

@@ -19,7 +19,7 @@ import java.util.UUID
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ThreadLocalRandom
 import java.util.concurrent.TimeUnit
-import java.util.function.Function
+import java.util.concurrent.atomic.AtomicReference
 import kotlin.math.min
 import kotlin.math.pow
 
@@ -84,52 +84,115 @@ private constructor(
             !modifiedRequest.headers.names().contains("X-Stainless-Retry-Count")
 
         var retries = 0
+        val result = CompletableFuture<HttpResponse>()
+        val currentResponseFuture = AtomicReference<CompletableFuture<HttpResponse>?>(null)
+        val currentSleepFuture = AtomicReference<CompletableFuture<Void>?>(null)
+
+        fun <T> trackFuture(
+            futureReference: AtomicReference<CompletableFuture<T>?>,
+            future: CompletableFuture<T>,
+        ) {
+            futureReference.set(future)
+            if (result.isCancelled) {
+                futureReference.getAndSet(null)?.cancel(false)
+            }
+        }
+
+        fun completeResult(response: HttpResponse?, throwable: Throwable?) {
+            if (throwable == null) {
+                val completedResponse = checkNotNull(response)
+                if (!result.complete(completedResponse)) {
+                    completedResponse.close()
+                }
+            } else {
+                result.completeExceptionally(throwable)
+            }
+        }
 
         fun executeWithRetries(
             request: HttpRequest,
             requestOptions: RequestOptions,
-        ): CompletableFuture<HttpResponse> {
+        ) {
+            if (result.isDone) {
+                return
+            }
+
             val requestWithRetryCount =
                 if (shouldSendRetryCount) setRetryCountHeader(request, retries) else request
 
             val responseFuture = httpClient.executeAsync(requestWithRetryCount, requestOptions)
+            trackFuture(currentResponseFuture, responseFuture)
             if (!isRetryable(requestWithRetryCount)) {
-                return responseFuture
+                responseFuture.whenComplete nonRetryableResponse@{ response, throwable ->
+                    currentResponseFuture.compareAndSet(responseFuture, null)
+
+                    if (result.isDone) {
+                        response?.close()
+                        return@nonRetryableResponse
+                    }
+
+                    completeResult(response, throwable)
+                }
+                return
             }
 
-            return responseFuture
-                .handleAsync(
-                    fun(
-                        response: HttpResponse?,
-                        throwable: Throwable?,
-                    ): CompletableFuture<HttpResponse> {
-                        if (response != null) {
-                            if (++retries > maxRetries || !shouldRetry(response)) {
-                                return CompletableFuture.completedFuture(response)
-                            }
-                        } else {
-                            if (++retries > maxRetries || !shouldRetry(throwable!!)) {
-                                val failedFuture = CompletableFuture<HttpResponse>()
-                                failedFuture.completeExceptionally(throwable)
-                                return failedFuture
-                            }
-                        }
+            responseFuture.whenComplete responseHandler@{ response, throwable ->
+                currentResponseFuture.compareAndSet(responseFuture, null)
 
-                        val backoffDuration = getRetryBackoffDuration(retries, response)
-                        // All responses must be closed, so close the failed one before retrying.
-                        response?.close()
-                        return sleeper.sleepAsync(backoffDuration).thenCompose {
-                            executeWithRetries(requestWithRetryCount, requestOptions)
-                        }
+                if (result.isDone) {
+                    response?.close()
+                    return@responseHandler
+                }
+
+                if (response != null) {
+                    if (++retries > maxRetries || !shouldRetry(response)) {
+                        completeResult(response, null)
+                        return@responseHandler
+                    }
+                } else {
+                    val actualThrowable = checkNotNull(throwable)
+                    if (++retries > maxRetries || !shouldRetry(actualThrowable)) {
+                        completeResult(null, actualThrowable)
+                        return@responseHandler
                     }
-                ) {
-                    // Run in the same thread.
-                    it.run()
                 }
-                .thenCompose(Function.identity())
+
+                val backoffDuration = getRetryBackoffDuration(retries, response)
+                // All responses must be closed, so close the failed one before retrying.
+                response?.close()
+
+                if (result.isDone) {
+                    return@responseHandler
+                }
+
+                val sleepFuture = sleeper.sleepAsync(backoffDuration)
+                trackFuture(currentSleepFuture, sleepFuture)
+                sleepFuture.whenComplete sleepHandler@{ _, sleepThrowable ->
+                    currentSleepFuture.compareAndSet(sleepFuture, null)
+
+                    if (result.isDone) {
+                        return@sleepHandler
+                    }
+
+                    if (sleepThrowable != null) {
+                        result.completeExceptionally(sleepThrowable)
+                        return@sleepHandler
+                    }
+
+                    executeWithRetries(requestWithRetryCount, requestOptions)
+                }
+            }
+        }
+
+        result.whenComplete { _, _ ->
+            if (result.isCancelled) {
+                currentResponseFuture.getAndSet(null)?.cancel(false)
+                currentSleepFuture.getAndSet(null)?.cancel(false)
+            }
         }
 
-        return executeWithRetries(modifiedRequest, requestOptions)
+        executeWithRetries(modifiedRequest, requestOptions)
+        return result
     }
 
     override fun close() {

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

@@ -19,6 +19,7 @@ import com.openai.client.okhttp.OkHttpClient
 import com.openai.core.RequestOptions
 import com.openai.core.Sleeper
 import com.openai.errors.OpenAIRetryableException
+import java.io.ByteArrayInputStream
 import java.io.InputStream
 import java.time.Clock
 import java.time.Duration
@@ -28,6 +29,7 @@ import java.time.format.DateTimeFormatter
 import java.util.concurrent.CompletableFuture
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.parallel.ResourceLock
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
@@ -40,21 +42,43 @@ internal class RetryingHttpClientTest {
     private lateinit var baseUrl: String
     private lateinit var httpClient: HttpClient
 
-    private class RecordingSleeper : Sleeper {
+    private class RecordingSleeper(
+        private val sleepAsyncFutureFactory: () -> CompletableFuture<Void> = {
+            CompletableFuture.completedFuture(null)
+        }
+    ) : Sleeper {
         val durations = mutableListOf<Duration>()
+        val sleepFutures = mutableListOf<CompletableFuture<Void>>()
 
         override fun sleep(duration: Duration) {
             durations.add(duration)
         }
 
         override fun sleepAsync(duration: Duration): CompletableFuture<Void> {
             durations.add(duration)
-            return CompletableFuture.completedFuture(null)
+            return sleepAsyncFutureFactory().also(sleepFutures::add)
         }
 
         override fun close() {}
     }
 
+    private class TestHttpResponse(
+        private val statusCode: Int,
+        private val headers: Headers = Headers.builder().build(),
+    ) : HttpResponse {
+        var isClosed = false
+
+        override fun statusCode(): Int = statusCode
+
+        override fun headers(): Headers = headers
+
+        override fun body(): InputStream = ByteArrayInputStream(byteArrayOf())
+
+        override fun close() {
+            isClosed = true
+        }
+    }
+
     @BeforeEach
     fun beforeEach(wmRuntimeInfo: WireMockRuntimeInfo) {
         baseUrl = wmRuntimeInfo.httpBaseUrl
@@ -124,6 +148,90 @@ internal class RetryingHttpClientTest {
         assertNoResponseLeaks()
     }
 
+    @Test
+    fun executeAsync_whenFutureCancelled_cancelsUnderlyingRequest() {
+        val responseFuture = CompletableFuture<HttpResponse>()
+        val retryingClient =
+            RetryingHttpClient.builder()
+                .httpClient(
+                    object : HttpClient {
+                        override fun execute(
+                            request: HttpRequest,
+                            requestOptions: RequestOptions,
+                        ): HttpResponse = throw UnsupportedOperationException()
+
+                        override fun executeAsync(
+                            request: HttpRequest,
+                            requestOptions: RequestOptions,
+                        ): CompletableFuture<HttpResponse> = responseFuture
+
+                        override fun close() {}
+                    }
+                )
+                .sleeper(RecordingSleeper())
+                .build()
+
+        val retriedResponseFuture =
+            retryingClient.executeAsync(
+                HttpRequest.builder()
+                    .method(HttpMethod.POST)
+                    .baseUrl(baseUrl)
+                    .addPathSegment("something")
+                    .build()
+            )
+
+        retriedResponseFuture.cancel(false)
+
+        assertThat(responseFuture).isCancelled()
+    }
+
+    @Test
+    fun executeAsync_whenFutureCancelledDuringBackoff_cancelsPendingSleepAndStopsRetrying() {
+        val attemptFutures = mutableListOf<CompletableFuture<HttpResponse>>()
+        val sleeper = RecordingSleeper { CompletableFuture<Void>() }
+        val retryingClient =
+            RetryingHttpClient.builder()
+                .httpClient(
+                    object : HttpClient {
+                        override fun execute(
+                            request: HttpRequest,
+                            requestOptions: RequestOptions,
+                        ): HttpResponse = throw UnsupportedOperationException()
+
+                        override fun executeAsync(
+                            request: HttpRequest,
+                            requestOptions: RequestOptions,
+                        ): CompletableFuture<HttpResponse> =
+                            CompletableFuture<HttpResponse>().also(attemptFutures::add)
+
+                        override fun close() {}
+                    }
+                )
+                .sleeper(sleeper)
+                .maxRetries(2)
+                .build()
+
+        val retriedResponseFuture =
+            retryingClient.executeAsync(
+                HttpRequest.builder()
+                    .method(HttpMethod.POST)
+                    .baseUrl(baseUrl)
+                    .addPathSegment("something")
+                    .build()
+            )
+        val failedResponse = TestHttpResponse(503)
+
+        attemptFutures.single().complete(failedResponse)
+
+        val sleepFuture = sleeper.sleepFutures.single()
+        retriedResponseFuture.cancel(false)
+        sleepFuture.complete(null)
+
+        assertThat(failedResponse.isClosed).isTrue()
+        assertThat(sleepFuture).isCancelled()
+        assertThat(attemptFutures).hasSize(1)
+    }
+
     @ParameterizedTest
     @ValueSource(booleans = [false, true])
     fun execute_withIdempotencyHeader(async: Boolean) {