use of io.envoyproxy.envoy.extensions.common.ratelimit.v3.RateLimitDescriptor in project Sentinel by alibaba.
the class SentinelEnvoyRlsServiceImpl method shouldRateLimit.
@Override
public void shouldRateLimit(RateLimitRequest request, StreamObserver<RateLimitResponse> responseObserver) {
int acquireCount = request.getHitsAddend();
if (acquireCount < 0) {
responseObserver.onError(new IllegalArgumentException("acquireCount should be positive, but actual: " + acquireCount));
return;
}
if (acquireCount == 0) {
// Not present, use the default "1" by default.
acquireCount = 1;
}
String domain = request.getDomain();
boolean blocked = false;
List<DescriptorStatus> statusList = new ArrayList<>(request.getDescriptorsCount());
for (RateLimitDescriptor descriptor : request.getDescriptorsList()) {
Tuple2<FlowRule, TokenResult> t = checkToken(domain, descriptor, acquireCount);
TokenResult r = t.r2;
printAccessLogIfNecessary(domain, descriptor, r);
if (r.getStatus() == TokenResultStatus.NO_RULE_EXISTS) {
// If the rule of the descriptor is absent, the request will pass directly.
r.setStatus(TokenResultStatus.OK);
}
if (!blocked && r.getStatus() != TokenResultStatus.OK) {
blocked = true;
}
Code statusCode = r.getStatus() == TokenResultStatus.OK ? Code.OK : Code.OVER_LIMIT;
DescriptorStatus.Builder descriptorStatusBuilder = DescriptorStatus.newBuilder().setCode(statusCode);
if (t.r1 != null) {
descriptorStatusBuilder.setCurrentLimit(RateLimit.newBuilder().setUnit(RateLimit.Unit.SECOND).setRequestsPerUnit((int) t.r1.getCount()).build()).setLimitRemaining(r.getRemaining());
}
statusList.add(descriptorStatusBuilder.build());
}
Code overallStatus = blocked ? Code.OVER_LIMIT : Code.OK;
RateLimitResponse response = RateLimitResponse.newBuilder().setOverallCode(overallStatus).addAllStatuses(statusList).build();
responseObserver.onNext(response);
responseObserver.onCompleted();
}
use of io.envoyproxy.envoy.extensions.common.ratelimit.v3.RateLimitDescriptor in project Sentinel by alibaba.
the class SentinelEnvoyRlsServiceImplTest method testShouldRateLimitPass.
@Test
public void testShouldRateLimitPass() {
SentinelEnvoyRlsServiceImpl rlsService = mock(SentinelEnvoyRlsServiceImpl.class);
StreamObserver<RateLimitResponse> streamObserver = mock(StreamObserver.class);
String domain = "testShouldRateLimitPass";
int acquireCount = 1;
RateLimitDescriptor descriptor1 = RateLimitDescriptor.newBuilder().addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk1").setValue("rv1").build()).build();
RateLimitDescriptor descriptor2 = RateLimitDescriptor.newBuilder().addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk2").setValue("rv2").build()).addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk3").setValue("rv3").build()).build();
ArgumentCaptor<RateLimitResponse> responseCapture = ArgumentCaptor.forClass(RateLimitResponse.class);
doNothing().when(streamObserver).onNext(responseCapture.capture());
doCallRealMethod().when(rlsService).shouldRateLimit(any(), any());
when(rlsService.checkToken(eq(domain), same(descriptor1), eq(acquireCount))).thenReturn(Tuple2.of(new FlowRule(), new TokenResult(TokenResultStatus.OK)));
when(rlsService.checkToken(eq(domain), same(descriptor2), eq(acquireCount))).thenReturn(Tuple2.of(new FlowRule(), new TokenResult(TokenResultStatus.OK)));
RateLimitRequest rateLimitRequest = RateLimitRequest.newBuilder().addDescriptors(descriptor1).addDescriptors(descriptor2).setDomain(domain).setHitsAddend(acquireCount).build();
rlsService.shouldRateLimit(rateLimitRequest, streamObserver);
RateLimitResponse response = responseCapture.getValue();
assertEquals(RateLimitResponse.Code.OK, response.getOverallCode());
response.getStatusesList().forEach(e -> assertEquals(RateLimitResponse.Code.OK, e.getCode()));
}
use of io.envoyproxy.envoy.extensions.common.ratelimit.v3.RateLimitDescriptor in project Sentinel by alibaba.
the class SentinelEnvoyRlsServiceImplTest method testShouldRatePartialBlock.
@Test
public void testShouldRatePartialBlock() {
SentinelEnvoyRlsServiceImpl rlsService = mock(SentinelEnvoyRlsServiceImpl.class);
StreamObserver<RateLimitResponse> streamObserver = mock(StreamObserver.class);
String domain = "testShouldRatePartialBlock";
int acquireCount = 1;
RateLimitDescriptor descriptor1 = RateLimitDescriptor.newBuilder().addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk1").setValue("rv1").build()).build();
RateLimitDescriptor descriptor2 = RateLimitDescriptor.newBuilder().addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk2").setValue("rv2").build()).addEntries(RateLimitDescriptor.Entry.newBuilder().setKey("rk3").setValue("rv3").build()).build();
ArgumentCaptor<RateLimitResponse> responseCapture = ArgumentCaptor.forClass(RateLimitResponse.class);
doNothing().when(streamObserver).onNext(responseCapture.capture());
doCallRealMethod().when(rlsService).shouldRateLimit(any(), any());
when(rlsService.checkToken(eq(domain), same(descriptor1), eq(acquireCount))).thenReturn(Tuple2.of(new FlowRule(), new TokenResult(TokenResultStatus.BLOCKED)));
when(rlsService.checkToken(eq(domain), same(descriptor2), eq(acquireCount))).thenReturn(Tuple2.of(new FlowRule(), new TokenResult(TokenResultStatus.OK)));
RateLimitRequest rateLimitRequest = RateLimitRequest.newBuilder().addDescriptors(descriptor1).addDescriptors(descriptor2).setDomain(domain).setHitsAddend(acquireCount).build();
rlsService.shouldRateLimit(rateLimitRequest, streamObserver);
RateLimitResponse response = responseCapture.getValue();
assertEquals(RateLimitResponse.Code.OVER_LIMIT, response.getOverallCode());
assertEquals(2, response.getStatusesCount());
assertTrue(response.getStatusesList().stream().anyMatch(e -> e.getCode().equals(RateLimitResponse.Code.OVER_LIMIT)));
assertFalse(response.getStatusesList().stream().allMatch(e -> e.getCode().equals(RateLimitResponse.Code.OVER_LIMIT)));
}
Aggregations