diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt index b9d7eef8f..925cc4abb 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt +++ b/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() + val currentResponseFuture = AtomicReference?>(null) + val currentSleepFuture = AtomicReference?>(null) + + fun trackFuture( + futureReference: AtomicReference?>, + future: CompletableFuture, + ) { + 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 { + ) { + 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 { - if (response != null) { - if (++retries > maxRetries || !shouldRetry(response)) { - return CompletableFuture.completedFuture(response) - } - } else { - if (++retries > maxRetries || !shouldRetry(throwable!!)) { - val failedFuture = CompletableFuture() - 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() { diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/RetryingHttpClientTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/RetryingHttpClientTest.kt index 21834efab..886c7bbaf 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/http/RetryingHttpClientTest.kt +++ b/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,8 +42,13 @@ internal class RetryingHttpClientTest { private lateinit var baseUrl: String private lateinit var httpClient: HttpClient - private class RecordingSleeper : Sleeper { + private class RecordingSleeper( + private val sleepAsyncFutureFactory: () -> CompletableFuture = { + CompletableFuture.completedFuture(null) + } + ) : Sleeper { val durations = mutableListOf() + val sleepFutures = mutableListOf>() override fun sleep(duration: Duration) { durations.add(duration) @@ -49,12 +56,29 @@ internal class RetryingHttpClientTest { override fun sleepAsync(duration: Duration): CompletableFuture { 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() + 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 = 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>() + val sleeper = RecordingSleeper { CompletableFuture() } + 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 = + CompletableFuture().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) {