fix: RouteController is not responsive #623 (#624)

epam · Dec 20, 2024 · d2039ad · d2039ad
1 parent 9ee09ec
commit d2039ad
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 15 deletions.
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
@@ -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
@@ -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
@@ -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
@@ -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<Arguments> datasource() {
         return List.of(
                 Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_user_key", 200, "/"),