Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Rate limit is not updated in RouteController #600 #601

Merged
merged 2 commits into from
Dec 5, 2024
Merged

fix: Rate limit is not updated in RouteController #600 #601

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fe6acd5
fix: Rate limit is not updated in RouteController #600
astsiapanay Dec 5, 2024
1064fb6
Merge branch 'development' into fix/issue-600
astsiapanay Dec 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,7 +415,7 @@ private Future<TokenUsage> collectTokenUsage(Buffer responseBody) {
tokenUsage = new TokenUsage();
}
context.setTokenUsage(tokenUsage);
proxy.getRateLimiter().increase(context).onFailure(error -> log.warn("Failed to increase limit. Trace: {}. Span: {}",
proxy.getRateLimiter().increase(context, context.getDeployment()).onFailure(error -> log.warn("Failed to increase limit. Trace: {}. Span: {}",
context.getTraceId(), context.getSpanId(), error));
try {
BigDecimal cost = ModelCostCalculator.calculate(context);
Expand Down
2 changes: 2 additions & 0 deletions server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,8 @@ private void handleProxyResponse(HttpClientResponse proxyResponse) {

if (responseStatusCode == 200) {
context.getUpstreamRoute().succeed();
proxy.getRateLimiter().increase(context, context.getRoute()).onFailure(error -> log.warn("Failed to increase limit. Trace: {}. Span: {}",
context.getTraceId(), context.getSpanId(), error));
}

BufferingReadStream proxyResponseStream = new BufferingReadStream(proxyResponse,
Expand Down
4 changes: 2 additions & 2 deletions server/src/main/java/com/epam/aidial/core/server/limiter/RateLimiter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public class RateLimiter {

private final ResourceService resourceService;

public Future<Void> increase(ProxyContext context) {
public Future<Void> increase(ProxyContext context, RoleBasedEntity roleBasedEntity) {
try {
// skip checking limits if redis is not available
if (resourceService == null) {
Expand All @@ -47,7 +47,7 @@ public Future<Void> increase(ProxyContext context) {
return Future.succeededFuture();
}

String tokensPath = getPathToTokens(context.getDeployment().getName());
String tokensPath = getPathToTokens(roleBasedEntity.getName());
ResourceDescriptor resourceDescription = getResourceDescription(context, tokensPath);
return vertx.executeBlocking(() -> updateTokenLimit(resourceDescription, usage.getTotalTokens()), false);
} catch (Throwable e) {
Expand Down
6 changes: 3 additions & 3 deletions ...er/src/test/java/com/epam/aidial/core/server/controller/DeploymentPostControllerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,13 +377,13 @@ public void testHandleResponse_Model() {
when(context.getUpstreamRoute()).thenReturn(upstreamRoute);
when(context.getResponseBody()).thenReturn(Buffer.buffer());
when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker);
when(rateLimiter.increase(any(ProxyContext.class))).thenReturn(Future.succeededFuture());
when(rateLimiter.increase(any(ProxyContext.class), eq(model))).thenReturn(Future.succeededFuture());
when(tokenStatsTracker.updateModelStats(context)).thenReturn(Future.succeededFuture());
BufferingReadStream bufferingReadStream = mock(BufferingReadStream.class);

controller.handleResponse(bufferingReadStream);

verify(rateLimiter).increase(eq(context));
verify(rateLimiter).increase(eq(context), eq(model));
verify(context).setTokenUsage(any(TokenUsage.class));
verify(logStore).save(eq(context));
verify(tokenStatsTracker).endSpan(eq(context));
Expand All @@ -409,7 +409,7 @@ public void testHandleResponse_App() {

controller.handleResponse(bufferingReadStream);

verify(rateLimiter, never()).increase(eq(context));
verify(rateLimiter, never()).increase(eq(context), eq(app));
verify(tokenStatsTracker).getTokenStats(eq(context));
verify(context).setTokenUsage(any(TokenUsage.class));
verify(logStore).save(eq(context));
Expand Down
20 changes: 10 additions & 10 deletions server/src/test/java/com/epam/aidial/core/server/limiter/RateLimiterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ public void testLimit_ApiKeySuccess_KeyExist() {
tokenUsage.setTotalTokens(90);
proxyContext.setTokenUsage(tokenUsage);

Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand All @@ -206,7 +206,7 @@ public void testLimit_ApiKeySuccess_KeyExist() {
assertNotNull(checkLimitFuture.result());
assertEquals(HttpStatus.OK, checkLimitFuture.result().status());

increaseLimitFuture = rateLimiter.increase(proxyContext);
increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand Down Expand Up @@ -254,7 +254,7 @@ public void testGetLimitStats_ApiKey() {
assertNotNull(resultFuture.result());
assertEquals(HttpStatus.OK, resultFuture.result().status());

Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand All @@ -272,7 +272,7 @@ public void testGetLimitStats_ApiKey() {
assertEquals(2, limitStats.getHourRequestStats().getTotal());
assertEquals(1, limitStats.getHourRequestStats().getUsed());

increaseLimitFuture = rateLimiter.increase(proxyContext);
increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand Down Expand Up @@ -325,7 +325,7 @@ public void testLimit_User_LimitFound() {
tokenUsage.setTotalTokens(150);
proxyContext.setTokenUsage(tokenUsage);

Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand All @@ -335,7 +335,7 @@ public void testLimit_User_LimitFound() {
assertNotNull(checkLimitFuture.result());
assertEquals(HttpStatus.OK, checkLimitFuture.result().status());

increaseLimitFuture = rateLimiter.increase(proxyContext);
increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand Down Expand Up @@ -367,7 +367,7 @@ public void testLimit_User_DefaultLimit() {
tokenUsage.setTotalTokens(90);
proxyContext.setTokenUsage(tokenUsage);

Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand All @@ -377,7 +377,7 @@ public void testLimit_User_DefaultLimit() {
assertNotNull(checkLimitFuture.result());
assertEquals(HttpStatus.OK, checkLimitFuture.result().status());

increaseLimitFuture = rateLimiter.increase(proxyContext);
increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand Down Expand Up @@ -423,7 +423,7 @@ public void testLimit_User_RequestLimit() {
tokenUsage.setTotalTokens(150);
proxyContext.setTokenUsage(tokenUsage);

Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand All @@ -433,7 +433,7 @@ public void testLimit_User_RequestLimit() {
assertNotNull(checkLimitFuture.result());
assertEquals(HttpStatus.OK, checkLimitFuture.result().status());

increaseLimitFuture = rateLimiter.increase(proxyContext);
increaseLimitFuture = rateLimiter.increase(proxyContext, model);
assertNotNull(increaseLimitFuture);
assertNull(increaseLimitFuture.cause());

Expand Down
Loading