RateLimitProfile.java

package se.jobtechdev.personaldatagateway.api.ratelimit;

import static java.time.Duration.ofSeconds;

import io.github.bucket4j.BandwidthBuilder;
import io.github.bucket4j.Bucket;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

public class RateLimitProfile {
  private final Map<String, Bucket> quotaByClient;
  private final String profileName;
  private final int maxTokens;
  private final int refillTokens;
  private final Duration refillPeriod;

  public record Quota(String rateLimitPolicy, String rateLimit, boolean exceeded) {}

  public RateLimitProfile(
      Map<String, Bucket> quotaByClient,
      String profileName,
      int maxTokens,
      int refillTokens,
      long refillPeriodInSeconds) {
    this.quotaByClient = quotaByClient;
    this.profileName = profileName;
    this.maxTokens = maxTokens;
    this.refillTokens = refillTokens;
    this.refillPeriod = ofSeconds(refillPeriodInSeconds);
  }

  protected static Function<
          BandwidthBuilder.BandwidthBuilderCapacityStage,
          BandwidthBuilder.BandwidthBuilderBuildStage>
      createLimitLambda(int maxTokens, int refillTokens, Duration refillPeriod) {
    return limit -> limit.capacity(maxTokens).refillIntervally(refillTokens, refillPeriod);
  }

  protected static String rateLimitPolicy(String profileName, long q, long w) {
    return String.format("\"%s\";q=%s;w=%s", profileName, q, w);
  }

  protected static String rateLimit(String profileName, long r, long t) {
    return String.format("\"%s\";r=%s;t=%s", profileName, r, t);
  }

  protected Bucket createQuota(String ignoredDummy) {
    return Bucket.builder()
        .addLimit(createLimitLambda(maxTokens, refillTokens, refillPeriod))
        .build();
  }

  public Quota consume(String key) {
    final var quota = quotaByClient.computeIfAbsent(profileName + "-" + key, this::createQuota);
    final var probe = quota.tryConsumeAndReturnRemaining(1);
    return new Quota(
        rateLimitPolicy(profileName, refillTokens, refillPeriod.toSeconds()),
        rateLimit(
            profileName,
            probe.getRemainingTokens(),
            TimeUnit.NANOSECONDS.toSeconds(probe.getNanosToWaitForReset())),
        !probe.isConsumed());
  }
}