TokenBucketAdmissionController implements AdmissionController { @Override public AdmissionControllerResponse apply(AdmissionControllerRequest request) { return findTokenBucket(request).map(this::consume).orElse(DEFAULT_OK); } TokenBucketAdmissionController(List<TokenBucketConfiguration> tokenBucketConfigurations,
                                          TitusRuntime titusRuntime); @Override AdmissionControllerResponse apply(AdmissionControllerRequest request);  }
@Test public void testSharedBucket() { TokenBucketAdmissionController controller = new TokenBucketAdmissionController( Collections.singletonList(SHARED_ANY_CONFIGURATION), titusRuntime ); AdmissionControllerRequest request = AdmissionControllerRequest.newBuilder() .withCallerId("any") .withEndpointName("any") .build(); int limit = SHARED_ANY_CONFIGURATION.getCapacity() + SHARED_ANY_CONFIGURATION.getRefillRateInSec() + 1; int stoppedAt = 0; while (stoppedAt < limit) { AdmissionControllerResponse response = controller.apply(request); if (!response.isAllowed()) { break; } stoppedAt++; } assertThat(stoppedAt).isGreaterThanOrEqualTo(SHARED_ANY_CONFIGURATION.getCapacity()); assertThat(stoppedAt).isLessThan(limit); }
@Test public void testNotSharedBucket() { TokenBucketAdmissionController controller = new TokenBucketAdmissionController( Collections.singletonList(NOT_SHARED_CONFIGURATION), titusRuntime ); AdmissionControllerRequest user1Request = AdmissionControllerRequest.newBuilder() .withCallerId("myUser1") .withEndpointName("any") .build(); AdmissionControllerRequest user2Request = AdmissionControllerRequest.newBuilder() .withCallerId("myUser2") .withEndpointName("any") .build(); int limit = NOT_SHARED_CONFIGURATION.getCapacity() + NOT_SHARED_CONFIGURATION.getRefillRateInSec() + 1; int stoppedAt = 0; while (stoppedAt < limit) { AdmissionControllerResponse response1 = controller.apply(user1Request); AdmissionControllerResponse response2 = controller.apply(user2Request); if (response1.isAllowed() && response2.isAllowed()) { stoppedAt++; } else { break; } } assertThat(stoppedAt).isGreaterThanOrEqualTo(NOT_SHARED_CONFIGURATION.getCapacity()); assertThat(stoppedAt).isLessThan(limit); }
@Test public void testOverlappingCallerIdButDifferentEndpointBuckets() { TokenBucketAdmissionController controller = new TokenBucketAdmissionController( Arrays.asList(SHARED_GETTERS_CONFIGURATION, SHARED_ANY_CONFIGURATION), titusRuntime ); AdmissionControllerRequest createRequest = AdmissionControllerRequest.newBuilder() .withCallerId("myUser") .withEndpointName("createX") .build(); AdmissionControllerRequest getRequest = AdmissionControllerRequest.newBuilder() .withCallerId("myUser") .withEndpointName("getX") .build(); AdmissionControllerResponse createResponse = controller.apply(createRequest); assertThat(createResponse.getReasonMessage()).contains(SHARED_ANY_CONFIGURATION.getName()); AdmissionControllerResponse getResponse = controller.apply(getRequest); assertThat(getResponse.getReasonMessage()).contains(SHARED_GETTERS_CONFIGURATION.getName()); }
@Test public void testNoMatch() { TokenBucketAdmissionController controller = new TokenBucketAdmissionController(Collections.emptyList(), titusRuntime); AdmissionControllerRequest request = AdmissionControllerRequest.newBuilder() .withCallerId("any") .withEndpointName("any") .build(); AdmissionControllerResponse response = controller.apply(request); assertThat(response.isAllowed()).isTrue(); assertThat(response.getReasonMessage()).isEqualTo("Rate limits not configured"); }