fix: log rate response request (#40)

Co-authored-by: Aliaksandr Stsiapanay <[email protected]>
epam · Nov 21, 2023 · 788e94e · 788e94e
1 parent f914774
commit 788e94e
Showing 1 changed file with 12 additions and 16 deletions.
diff --git a/src/main/java/com/epam/aidial/core/controller/RateResponseController.java b/src/main/java/com/epam/aidial/core/controller/RateResponseController.java
@@ -38,15 +38,10 @@ public void handle(String deploymentId) {
             return;
         }
 
-        if (deployment.getRateEndpoint() == null) {
-            context.respond(HttpStatus.OK);
-            proxy.getLogStore().save(context);
-        } else {
-            context.setDeployment(deployment);
-            context.getRequest().body()
-                    .onSuccess(this::handleRequestBody)
-                    .onFailure(this::handleRequestBodyError);
-        }
+        context.setDeployment(deployment);
+        context.getRequest().body()
+                .onSuccess(this::handleRequestBody)
+                .onFailure(this::handleRequestBodyError);
     }
 
     private Deployment getDeployment(String id) {
@@ -81,7 +76,13 @@ private void sendRequest() {
 
     private void handleRequestBody(Buffer requestBody) {
         context.setRequestBody(requestBody);
-        sendRequest();
+        Deployment deployment = context.getDeployment();
+        if (deployment.getRateEndpoint() == null) {
+            context.respond(HttpStatus.OK);
+            proxy.getLogStore().save(context);
+        } else {
+            sendRequest();
+        }
     }
 
     /**
@@ -110,11 +111,6 @@ private void handleProxyResponse(HttpClientResponse proxyResponse) {
         log.info("Received response header from origin: status={}, headers={}", proxyResponse.statusCode(),
                 proxyResponse.headers().size());
 
-        if (proxyResponse.statusCode() == HttpStatus.TOO_MANY_REQUESTS.getCode()) {
-            sendRequest(); // try next
-            return;
-        }
-
         BufferingReadStream proxyResponseStream = new BufferingReadStream(proxyResponse,
                 ProxyUtil.contentLength(proxyResponse, 1024));
 
@@ -155,7 +151,7 @@ private void handleRequestBodyError(Throwable error) {
      */
     private void handleProxyConnectionError(Throwable error) {
         log.warn("Can't connect to origin: {}", error.getMessage());
-        sendRequest(); // try next
+        context.respond(HttpStatus.BAD_GATEWAY, "connection error to origin");
     }
 
     /**