io.stargate.db.limiter.RateLimitingDecision Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of persistence-api Show documentation
There is a newer version: 2.1.0-BETA-19
package io.stargate.db.limiter;

import java.util.concurrent.CompletableFuture;
import java.util.function.Supplier;
import org.apache.cassandra.stargate.exceptions.UnauthorizedException;

/**
 * The decision taken by a {@link RateLimitingManager} for a particular query.
 *
 * The is essentially 3 possible decision:
 *
 * 

 *   to not rate limit at all ({@link #unlimited()}).
 *   
to rate limit the query, using a provided limiter ({@link #limit}).
 *   
to reject the query altogether ({@link #reject}).
 * 
 */
public abstract class RateLimitingDecision {

  private RateLimitingDecision() {}

  /** Creates a new decision consisting of not rate limiting a query. */
  public static Unlimited unlimited() {
    return Unlimited.INSTANCE;
  }

  /**
   * Creates a new decision consisting of rate limiting a query through acquiring the provided
   * number of permit on the provided limiter.
   */
  public static Limited limit(AsyncRateLimiter limiter, long permitsToAcquire) {
    return new Limited(limiter, permitsToAcquire);
  }

  /**
   * Creates a new decision consisting of rejecting a query, the rejected query throwing an {@link
   * UnauthorizedException} with the provided message.
   */
  public static Rejected reject(String rejectionMessage) {
    return new Rejected(rejectionMessage);
  }

  /** Applies this decision to the provided asynchronous taks/query. */
  public abstract  CompletableFuture apply(Supplier> task);

  public static class Unlimited extends RateLimitingDecision {
    private static final Unlimited INSTANCE = new Unlimited();

    @Override
    public  CompletableFuture apply(Supplier> task) {
      return task.get();
    }
  }

  public static class Limited extends RateLimitingDecision {
    private final AsyncRateLimiter limiter;
    private final long permitsToAcquire;

    private Limited(AsyncRateLimiter limiter, long permitsToAcquire) {
      this.limiter = limiter;
      this.permitsToAcquire = permitsToAcquire;
    }

    @Override
    public  CompletableFuture apply(Supplier> task) {
      return limiter.acquireAndExecute(permitsToAcquire, task);
    }
  }

  public static class Rejected extends RateLimitingDecision {
    private final String rejectionMessage;

    private Rejected(String rejectionMessage) {
      this.rejectionMessage = rejectionMessage;
    }

    @Override
    public  CompletableFuture apply(Supplier> task) {
      CompletableFuture exceptionalFuture = new CompletableFuture<>();
      exceptionalFuture.completeExceptionally(new UnauthorizedException(rejectionMessage));
      return exceptionalFuture;
    }
  }
}