diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java index ba53a4af..45e31081 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java @@ -344,7 +344,9 @@ private void handleProxyResponse(HttpClientResponse proxyResponse) { if (responseStatusCode == 200) { upstreamRoute.succeed(); - } else if (isFailedStatusCode(responseStatusCode)) { + } else { + // mark the upstream as failed + // and the next time we will select another one upstreamRoute.fail(proxyResponse); } @@ -377,10 +379,6 @@ private boolean isRetriableError(int statusCode) { return DEFAULT_RETRIABLE_HTTP_CODES.contains(statusCode) || context.getConfig().getRetriableErrorCodes().contains(statusCode); } - private static boolean isFailedStatusCode(int statusCode) { - return statusCode == 429 || statusCode >= 500; - } - /** * Called when proxy sent response from the origin to the client. */ diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java b/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java index aab2e82b..22ef61da 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java @@ -161,8 +161,6 @@ private void handleProxyResponse(HttpClientResponse proxyResponse) { context.getUpstreamRoute().succeed(); proxy.getRateLimiter().increase(context, context.getRoute()).onFailure(error -> log.warn("Failed to increase limit. Trace: {}. Span: {}", context.getTraceId(), context.getSpanId(), error)); - } else { - context.getUpstreamRoute().fail(proxyResponse); } BufferingReadStream proxyResponseStream = new BufferingReadStream(proxyResponse, diff --git a/server/src/main/java/com/epam/aidial/core/server/upstream/UpstreamState.java b/server/src/main/java/com/epam/aidial/core/server/upstream/UpstreamState.java index f74ee1e2..03c7e31a 100644 --- a/server/src/main/java/com/epam/aidial/core/server/upstream/UpstreamState.java +++ b/server/src/main/java/com/epam/aidial/core/server/upstream/UpstreamState.java @@ -39,19 +39,15 @@ class UpstreamState { } /** - * Register upstream failure. Supported error codes are 429 and 5xx. + * Register upstream failure. * * @param status response status code from upstream - * @param retryAfterSeconds time in seconds when upstream may become available; only take into account with 429 status code + * @param retryAfterSeconds time in seconds when upstream may become available */ void fail(HttpStatus status, long retryAfterSeconds) { this.source = retryAfterSeconds == -1 ? RetryAfterSource.CORE : RetryAfterSource.UPSTREAM; this.status = status; - if (status == HttpStatus.TOO_MANY_REQUESTS) { - retryAfterSeconds = source == RetryAfterSource.CORE ? DEFAULT_RETRY_AFTER_SECONDS_VALUE : retryAfterSeconds; - setReplyAfter(retryAfterSeconds); - log.warn("Upstream {} limit hit: retry after {}", upstream.getEndpoint(), Instant.ofEpochMilli(retryAfter).toString()); - } else if (status.is5xx()) { + if (status.is5xx()) { if (source == RetryAfterSource.CORE) { if (errorCount != 30) { errorCount++; @@ -62,7 +58,11 @@ void fail(HttpStatus status, long retryAfterSeconds) { setReplyAfter(retryAfterSeconds); } } else { - throw new IllegalArgumentException("Unsupported http status: " + status); + retryAfterSeconds = source == RetryAfterSource.CORE ? DEFAULT_RETRY_AFTER_SECONDS_VALUE : retryAfterSeconds; + setReplyAfter(retryAfterSeconds); + if (status == HttpStatus.TOO_MANY_REQUESTS) { + log.warn("Upstream {} limit hit: retry after {}", upstream.getEndpoint(), Instant.ofEpochMilli(retryAfter).toString()); + } } } diff --git a/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java b/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java index d48d8e45..21f35dfe 100644 --- a/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java @@ -57,6 +57,15 @@ void routeNotRateLimited() { } } + @Test + void route_404() { + TestWebServer.Handler handler = request -> new MockResponse().setResponseCode(404); + try (TestWebServer server = new TestWebServer(9876, handler)) { + Response resp = send(HttpMethod.GET, "/v1/plain", null, null, "api-key", "vstore_user_key"); + assertEquals(404, resp.status()); + } + } + private static List datasource() { return List.of( Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_user_key", 200, "/"),