private synchronized void addPermitAllocation(PermitAllocation allocation) { this.batches.put(allocation.getExpiration(), new PermitBatch(allocation.getPermits(), allocation.getExpiration())); this.totalAvailablePermits += allocation.getPermits(); } }
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { long permitsRequested = request.getPermits(); Long minPermits = request.getMinPermits(GetMode.NULL); if (minPermits == null) { minPermits = permitsRequested; } long permitsGranted = this.tokenBucket.getPermits(permitsRequested, minPermits, LimiterServerResource.TIMEOUT_MILLIS); PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(permitsGranted); allocation.setExpiration(Long.MAX_VALUE); if (permitsGranted <= 0) { allocation.setMinRetryDelayMillis(LimiterServerResource.TIMEOUT_MILLIS); } return allocation; }
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { PermitAllocation allocation = new PermitAllocation(); // For overflow safety, don't return max long allocation.setPermits(Math.max(Long.MAX_VALUE / 100, request.getPermits())); allocation.setExpiration(Long.MAX_VALUE); return allocation; }
allocation = policy.computePermitAllocation(request); permitsGrantedMeter.mark(allocation.getPermits());
@Override public void onSuccess(Response<PermitAllocation> result) { BatchedPermitsRequester.this.retries = 0; BatchedPermitsRequester.this.callbackCounter.incrementAndGet(); BatchedPermitsRequester.this.lock.lock(); try { PermitAllocation allocation = result.getEntity(); log.debug("Received permit allocation " + allocation); Long retryDelay = allocation.getMinRetryDelayMillis(GetMode.NULL); if (retryDelay != null) { BatchedPermitsRequester.this.retryStatus.blockRetries(retryDelay, null); } if (allocation.getPermits() > 0) { BatchedPermitsRequester.this.permitBatchContainer.addPermitAllocation(allocation); } BatchedPermitsRequester.this.requestSemaphore.release(); if (allocation.getPermits() > 0) { BatchedPermitsRequester.this.newPermitsAvailable.signalAll(); } } finally { try { this.timerContext.close(); } catch (IOException ioe) { // Do nothing } BatchedPermitsRequester.this.lock.unlock(); } }
@Override public void sendRequest(PermitRequest request, Callback<Response<PermitAllocation>> callback) { this.requestList.add(request); PermitAllocation permitAllocation = new PermitAllocation(); permitAllocation.setPermits(request.getPermits()); permitAllocation.setExpiration(Long.MAX_VALUE); Response<PermitAllocation> response = Mockito.mock(Response.class); Mockito.when(response.getEntity()).thenReturn(permitAllocation); callback.onSuccess(response); } }
@Test public void test() throws Exception { try(Closer closer = Closer.create()) { Map<String, String> configMap = Maps.newHashMap(); TestingServer zkTestingServer = closer.register(new TestingServer(-1)); configMap.put(ThrottlingGuiceServletConfig.ZK_STRING_KEY, zkTestingServer.getConnectString()); configMap.put(ThrottlingGuiceServletConfig.HA_CLUSTER_NAME, TestFailover.class.getSimpleName() + "_cluster"); Config config = ConfigFactory.parseMap(configMap); ThrottlingGuiceServletConfig server2001 = createServerAtPort(config, 2001); PermitAllocation allocation = sendRequestToServer(server2001, 10); Assert.assertTrue(allocation.getPermits() >= 1); ThrottlingGuiceServletConfig server2002 = createServerAtPort(config, 2002); allocation = sendRequestToServer(server2001, 10); Assert.assertTrue(allocation.getPermits() >= 1); try { sendRequestToServer(server2002, 10); Assert.fail(); } catch (RestLiServiceException exc) { Assert.assertTrue(exc.hasErrorDetails()); Assert.assertTrue(exc.getErrorDetails().containsKey(LimiterServerResource.LOCATION_301)); Assert.assertEquals(new URI(exc.getErrorDetails().get(LimiterServerResource.LOCATION_301).toString()).getPort(), 2001); } server2001.close(); allocation = sendRequestToServer(server2002, 10); Assert.assertTrue(allocation.getPermits() >= 1); } }
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { long permits = request.getPermits(); long allocated = 0; try { if (limiter.acquirePermits(permits) != null) { allocated = permits; } else { throw new RestLiServiceException(HttpStatus.S_403_FORBIDDEN, "Not enough permits."); } } catch (InterruptedException ie) { // return no permits } PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(allocated); allocation.setExpiration(Long.MAX_VALUE); if (allocated <= 0) { allocation.setMinRetryDelayMillis(60000); } return allocation; }
@Builder(builderMethodName = "satisfyRequestBuilder", buildMethodName = "satisfy") public static void satisfyRequest(RequestAndCallback requestAndCallback, long expiration) { PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(requestAndCallback.getRequest().getPermits()); allocation.setExpiration(expiration > 0 ? expiration : Long.MAX_VALUE); Response<PermitAllocation> response = Mockito.mock(Response.class); Mockito.when(response.getEntity()).thenReturn(allocation); requestAndCallback.getCallback().onSuccess(response); }
@Test public void test() { ThrottlingGuiceServletConfig guiceServletConfig = new ThrottlingGuiceServletConfig(); guiceServletConfig.initialize(ConfigFactory.empty()); Injector injector = guiceServletConfig.getInjector(); LimiterServerResource limiterServer = injector.getInstance(LimiterServerResource.class); PermitRequest request = new PermitRequest(); request.setPermits(10); request.setResource("myResource"); PermitAllocation allocation = limiterServer.getSync(new ComplexResourceKey<>(request, new EmptyRecord())); Assert.assertTrue(allocation.getPermits() >= 10); }
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { long permitsRequested = request.getPermits(); Long minPermits = request.getMinPermits(GetMode.NULL); if (minPermits == null) { minPermits = permitsRequested; } long permitsGranted = this.tokenBucket.getPermits(permitsRequested, minPermits, LimiterServerResource.TIMEOUT_MILLIS); PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(permitsGranted); allocation.setExpiration(Long.MAX_VALUE); if (permitsGranted <= 0) { allocation.setMinRetryDelayMillis(LimiterServerResource.TIMEOUT_MILLIS); } return allocation; }
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { PermitAllocation allocation = new PermitAllocation(); // For overflow safety, don't return max long allocation.setPermits(Math.max(Long.MAX_VALUE / 100, request.getPermits())); allocation.setExpiration(Long.MAX_VALUE); return allocation; }
res3request.setResource("res3"); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(), new Long(20)); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(), new Long(20)); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(), new Long(20)); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(), new Long(20)); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(), new Long(20)); limiterServer.getSync(new ComplexResourceKey<>(res1request, new EmptyRecord())).getPermits(); Assert.fail(); } catch (RestLiServiceException exc) { Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res2request, new EmptyRecord())).getPermits(), new Long(20)); Assert.assertEquals(limiterServer.getSync(new ComplexResourceKey<>(res2request, new EmptyRecord())).getPermits(), new Long(20)); limiterServer.getSync(new ComplexResourceKey<>(res2request, new EmptyRecord())).getPermits(); Assert.fail(); } catch (RestLiServiceException exc) { Assert.assertTrue(limiterServer.getSync(new ComplexResourceKey<>(res3request, new EmptyRecord())).getPermits() >= res3request.getPermits());
@Override public PermitAllocation computePermitAllocation(PermitRequest request) { long permits = request.getPermits(); long allocated = 0; try { if (limiter.acquirePermits(permits) != null) { allocated = permits; } else { throw new RestLiServiceException(HttpStatus.S_403_FORBIDDEN, "Not enough permits."); } } catch (InterruptedException ie) { // return no permits } PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(allocated); allocation.setExpiration(Long.MAX_VALUE); if (allocated <= 0) { allocation.setMinRetryDelayMillis(60000); } return allocation; }
Assert.assertEquals(allocation.getPermits(), new Long(20)); Assert.assertEquals(allocation.getPermits(), new Long(20));
allocation = policy.computePermitAllocation(request); permitsGrantedMeter.mark(allocation.getPermits());