InMemoryRateLimitService extends AbstractRateLimitService { @Override public RateLimitDescriptor consume(HttpServletRequest request, long consumption) { synchronized (this) { RateLimitDescriptor descriptor = computeRateLimitRecovery(request); if (descriptor == null) { return null; } descriptor = specs.computeIfAbsent(descriptor.getLimitationUnitName(), p -> computeRateLimitRecovery(request)); if (descriptor == null) { return null; } long now = Clock.now().toEpochMillisec(); long secSinceLastUpdate = now - descriptor.getLastUpdateTime(); log.debug("Time (sec) since last update = {}", secSinceLastUpdate); long fill = secSinceLastUpdate * descriptor.getFillRate(); long budget = Math.min(descriptor.getMaxBudget(), descriptor.getCurrentBudget() + fill); log.info("Budget before current request (filled {}): {}", fill, budget); budget -= consumption; log.info("Budget after current request (consumed {}): {}", consumption, budget); descriptor.setLastUpdateTime(now); descriptor.setCurrentBudget(budget); return descriptor; } }  @Override RateLimitDescriptor consume(HttpServletRequest request, long consumption); @Override RateLimitDescriptor get(HttpServletRequest request);  }
@Test public void consume100() { RateLimitDescriptor actual = sut.consume(request, 100); assertThat(actual.getMaxBudget(), is(1000L)); assertThat(actual.getFillRate(), is(2L)); assertThat(actual.getCurrentBudget(), is(900L)); }
@Test public void consume100_consume200() { Clock.setTimeSource(new FixedTimeSource(TimePoint.EPOCH)); sut.consume(request, 100); RateLimitDescriptor actual = sut.consume(request, 200); assertThat(actual.getMaxBudget(), is(1000L)); assertThat(actual.getFillRate(), is(2L)); assertThat(actual.getCurrentBudget(), is(700L)); }
@Test public void consume100_recover20_consume200() { Clock.setTimeSource(new FixedTimeSource(TimePoint.EPOCH)); sut.consume(request, 100); Clock.setTimeSource(new FixedTimeSource(TimePoint.from(10L))); RateLimitDescriptor actual = sut.consume(request, 200); assertThat(actual.getMaxBudget(), is(1000L)); assertThat(actual.getFillRate(), is(2L)); assertThat(actual.getCurrentBudget(), is(720L)); }
@Test public void consume100_recover400_consume200() { Clock.setTimeSource(new FixedTimeSource(TimePoint.EPOCH)); sut.consume(request, 100); Clock.setTimeSource(new FixedTimeSource(TimePoint.from(200L))); RateLimitDescriptor actual = sut.consume(request, 200); assertThat(actual.getMaxBudget(), is(1000L)); assertThat(actual.getFillRate(), is(2L)); assertThat(actual.getCurrentBudget(), is(800L)); }