fix: Core rejects requests if trace id is not found #110 (#111)

Co-authored-by: Aliaksandr Stsiapanay <[email protected]>
epam · Jan 5, 2024 · e02bab8 · e02bab8
1 parent 6c2cb9b
commit e02bab8
Show file tree

Hide file tree

Showing 7 changed files with 84 additions and 244 deletions.
diff --git a/src/main/java/com/epam/aidial/core/Proxy.java b/src/main/java/com/epam/aidial/core/Proxy.java
@@ -14,6 +14,7 @@
 import com.epam.aidial.core.upstream.UpstreamBalancer;
 import com.epam.aidial.core.util.HttpStatus;
 import com.epam.aidial.core.util.ProxyUtil;
+import io.opentelemetry.api.trace.Span;
 import io.vertx.core.Future;
 import io.vertx.core.Handler;
 import io.vertx.core.Vertx;
@@ -115,6 +116,7 @@ private void handleRequest(HttpServerRequest request) {
         Config config = configStore.load();
         String apiKey = request.headers().get(HEADER_API_KEY);
         String authorization = request.getHeader(HttpHeaders.AUTHORIZATION);
+        String traceId = Span.current().getSpanContext().getTraceId();
         log.debug("Authorization header: {}", authorization);
         Key key;
         if (apiKey == null && authorization == null) {
@@ -143,7 +145,7 @@ private void handleRequest(HttpServerRequest request) {
         extractedClaims.onComplete(result -> {
             try {
                 if (result.succeeded()) {
-                    onExtractClaimsSuccess(result.result(), config, request, key);
+                    onExtractClaimsSuccess(result.result(), config, request, key, traceId);
                 } else {
                     onExtractClaimsFailure(result.cause(), request);
                 }
@@ -161,8 +163,8 @@ private void onExtractClaimsFailure(Throwable error, HttpServerRequest request)
     }
 
     private void onExtractClaimsSuccess(ExtractedClaims extractedClaims, Config config,
-                                        HttpServerRequest request, Key key) throws Exception {
-        ProxyContext context = new ProxyContext(config, request, key, extractedClaims);
+                                        HttpServerRequest request, Key key, String traceId) throws Exception {
+        ProxyContext context = new ProxyContext(config, request, key, extractedClaims, traceId);
         Controller controller = ControllerSelector.select(this, context);
         controller.handle();
     }

diff --git a/src/main/java/com/epam/aidial/core/ProxyContext.java b/src/main/java/com/epam/aidial/core/ProxyContext.java
@@ -31,6 +31,7 @@ public class ProxyContext {
     private final Key key;
     private final HttpServerRequest request;
     private final HttpServerResponse response;
+    private final String traceId;
 
     private Deployment deployment;
     private String userSub;
@@ -52,7 +53,7 @@ public class ProxyContext {
     // the project belongs to API key which initiated request
     private String originalProject;
 
-    public ProxyContext(Config config, HttpServerRequest request, Key key, ExtractedClaims extractedClaims) {
+    public ProxyContext(Config config, HttpServerRequest request, Key key, ExtractedClaims extractedClaims, String traceId) {
         this.config = config;
         this.key = key;
         if (key != null) {
@@ -67,6 +68,7 @@ public ProxyContext(Config config, HttpServerRequest request, Key key, Extracted
             this.userHash = extractedClaims.userHash();
             this.userSub = extractedClaims.sub();
         }
+        this.traceId = traceId;
     }
 
     public Future<Void> respond(HttpStatus status) {

diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java
@@ -57,6 +57,7 @@ public class DeploymentPostController {
 
     private final Proxy proxy;
     private final ProxyContext context;
+    private boolean unregisterTrace;
 
     public Future<?> handle(String deploymentId, String deploymentApi) {
         String contentType = context.getRequest().getHeader(HttpHeaders.CONTENT_TYPE);
@@ -79,10 +80,7 @@ public Future<?> handle(String deploymentId, String deploymentApi) {
             return context.respond(HttpStatus.FORBIDDEN, "Forbidden deployment");
         }
 
-        if (!proxy.getRateLimiter().register(context)) {
-            log.warn("Trace is not found by id for request: method={}, uri={}", context.getRequest().method(), context.getRequest().uri());
-            return respond(HttpStatus.BAD_REQUEST, "Trace is not found");
-        }
+        unregisterTrace = !proxy.getRateLimiter().register(context);
 
         context.setDeployment(deployment);
 
@@ -247,7 +245,7 @@ private void handleProxyResponse(HttpClientResponse proxyResponse) {
                 .to(response)
                 .onSuccess(ignored -> handleResponse())
                 .onFailure(this::handleResponseError)
-                .onComplete(ignore -> proxy.getRateLimiter().unregister());
+                .onComplete(ignore -> unregister());
     }
 
     /**
@@ -463,17 +461,23 @@ private static boolean isBaseAssistant(Deployment deployment) {
     }
 
     private Future<Void> respond(HttpStatus status, String errorMessage) {
-        proxy.getRateLimiter().unregister();
+        unregister();
         return context.respond(status, errorMessage);
     }
 
     private Future<Void> respond(HttpStatus status) {
-        proxy.getRateLimiter().unregister();
+        unregister();
         return context.respond(status);
     }
 
     private Future<Void> respond(HttpStatus status, Object result) {
-        proxy.getRateLimiter().unregister();
+        unregister();
         return context.respond(status, result);
     }
+
+    private void unregister() {
+        if (unregisterTrace) {
+            proxy.getRateLimiter().unregister(context);
+        }
+    }
 }
diff --git a/src/main/java/com/epam/aidial/core/limiter/RateLimiter.java b/src/main/java/com/epam/aidial/core/limiter/RateLimiter.java
@@ -8,7 +8,6 @@
 import com.epam.aidial.core.token.TokenUsage;
 import com.epam.aidial.core.util.HttpStatus;
 import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.sdk.trace.ReadableSpan;
 import lombok.extern.slf4j.Slf4j;
 
 import java.util.List;
@@ -20,7 +19,7 @@ public class RateLimiter {
     private final ConcurrentHashMap<Id, RateLimit> rates = new ConcurrentHashMap<>();
 
     public void increase(ProxyContext context) {
-        Entity entity = getEntityFromTracingContext();
+        Entity entity = getEntityFromTracingContext(context);
         if (entity == null || entity.user()) {
             return;
         }
@@ -39,7 +38,7 @@ public void increase(ProxyContext context) {
     }
 
     public RateLimitResult limit(ProxyContext context) {
-        Entity entity = getEntityFromTracingContext();
+        Entity entity = getEntityFromTracingContext(context);
         if (entity == null) {
             Span span = Span.current();
             log.warn("Entity is not found by traceId={}", span.getSpanContext().getTraceId());
@@ -75,36 +74,33 @@ public RateLimitResult limit(ProxyContext context) {
         return rate.update(timestamp, limit);
     }
 
+    /**
+     * Returns <code>true</code> if the trace is already registered otherwise <code>false</code>.
+     */
     public boolean register(ProxyContext context) {
-        ReadableSpan span = (ReadableSpan) Span.current();
-        String traceId = span.getSpanContext().getTraceId();
-        if (span.getParentSpanContext().isRemote()) {
-            Entity entity = traceIdToEntity.get(traceId);
-            if (entity != null) {
-                if (entity.user()) {
-                    context.setUserHash(entity.name());
-                } else {
-                    context.setOriginalProject(entity.name());
-                }
+        String traceId = context.getTraceId();
+        Entity entity = traceIdToEntity.get(traceId);
+        if (entity != null) {
+            // update context with the original requester
+            if (entity.user()) {
+                context.setUserHash(entity.name());
+            } else {
+                context.setOriginalProject(entity.name());
             }
-            return entity != null;
         } else {
             if (context.getKey() != null) {
                 Key key = context.getKey();
                 traceIdToEntity.put(traceId, new Entity(key.getKey(), List.of(key.getRole()), key.getProject(), false));
             } else {
                 traceIdToEntity.put(traceId, new Entity(context.getUserSub(), context.getUserRoles(), context.getUserHash(), true));
             }
-            return true;
         }
+        return entity != null;
     }
 
-    public void unregister() {
-        ReadableSpan span = (ReadableSpan) Span.current();
-        if (!span.getParentSpanContext().isRemote()) {
-            String traceId = span.getSpanContext().getTraceId();
-            traceIdToEntity.remove(traceId);
-        }
+    public void unregister(ProxyContext context) {
+        String traceId = context.getTraceId();
+        traceIdToEntity.remove(traceId);
     }
 
     private Limit getLimitByApiKey(ProxyContext context, Entity entity) {
@@ -120,9 +116,8 @@ private Limit getLimitByApiKey(ProxyContext context, Entity entity) {
         return role.getLimits().get(deployment.getName());
     }
 
-    protected Entity getEntityFromTracingContext() {
-        ReadableSpan span = (ReadableSpan) Span.current();
-        String traceId = span.getSpanContext().getTraceId();
+    private Entity getEntityFromTracingContext(ProxyContext context) {
+        String traceId = context.getTraceId();
         return traceIdToEntity.get(traceId);
     }
 

diff --git a/src/test/java/com/epam/aidial/core/TracerApiTest.java b/src/test/java/com/epam/aidial/core/TracerApiTest.java
diff --git a/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java b/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java
@@ -74,7 +74,6 @@ public void beforeEach() {
     @Test
     public void testUnsupportedContentType() {
         when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn("unsupported");
-        when(proxy.getRateLimiter()).thenReturn(rateLimiter);
 
         controller.handle("app1", "api");
 
@@ -112,23 +111,6 @@ public void testDeploymentNotFound() {
         verify(context).respond(eq(NOT_FOUND), anyString());
     }
 
-    @Test
-    public void testTraceNotFound() {
-        when(proxy.getRateLimiter()).thenReturn(rateLimiter);
-        when(rateLimiter.register(any(ProxyContext.class))).thenReturn(false);
-        when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn(HEADER_CONTENT_TYPE_APPLICATION_JSON);
-        Config config = new Config();
-        config.setApplications(new HashMap<>());
-        Application application = new Application();
-        application.setName("app1");
-        config.getApplications().put("app1", application);
-        when(context.getConfig()).thenReturn(config);
-
-        controller.handle("app1", "chat/completions");
-
-        verify(context).respond(eq(BAD_REQUEST), anyString());
-    }
-
     @Test
     public void testNoRoute() {
         when(proxy.getRateLimiter()).thenReturn(rateLimiter);