discord4j.rest.request.GlobalRateLimiter Maven / Gradle / Ivy

Go to download
/*
 * This file is part of Discord4J.
 *
 * Discord4J is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Discord4J is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with Discord4J.  If not, see .
 */
package discord4j.rest.request;

import reactor.core.Exceptions;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;

/**
 * Used to prevent requests from being sent while the bot is
 * globally rate limited.
 * 
 * Provides a single resource that can be acquired through the use of {@link #withLimiter(Supplier)}, blocking all other
 * attempts until the {@link Mono} supplier completes or terminates with an error.
 * 

 * This rate limiter can have their delay directly indicated through {@link #rateLimitFor(Duration)}, determining the
 * duration a resource holder must wait before processing starts.
 * 
 * When subscribing to the rate limiter, the only guarantee is that the subscription will be completed at some point in
 * the future. If a global ratelimit is in effect, it will be completed when the cooldown ends. Otherwise, it is
 * completed immediately.
 */
public class GlobalRateLimiter {

    private final Semaphore outer = new Semaphore(8, true);
    private final Semaphore inner = new Semaphore(1, true);
    private final AtomicLong limitedUntil = new AtomicLong(0L);

    /**
     * Sets a new rate limit that will be applied to every new resource acquired.
     *
     * @param duration the {@link Duration} every new acquired resource should wait before being used
     */
    public void rateLimitFor(Duration duration) {
        limitedUntil.set(System.nanoTime() + duration.toNanos());
    }

    /**
     * Returns a {@link Mono} indicating that the rate limit has ended.
     *
     * @return a {@link Mono} that completes when the currently set limit has completed
     */
    Mono onComplete() {
        return Mono.defer(this::notifier);
    }

    private Mono notifier() {
        long delayNanos = delayNanos();
        if (delayNanos > 0) {
            return Mono.delay(Duration.ofNanos(delayNanos)).then();
        }
        return Mono.empty();
    }

    private long delayNanos() {
        return limitedUntil.get() - System.nanoTime();
    }

    /**
     * Provides a scope to perform reactive operations under this limiter resources. Resources are acquired on
     * subscription and released when the given stage has completed or terminated with an error.
     *
     * @param stage a supplier containing a {@link Mono} that will manage this limiter resources
     * @param  the type of the stage supplier
     * @return a {@link Mono} where each subscription represents acquiring a rate limiter resource
     */
    public  Mono withLimiter(Supplier> stage) {
        return Mono.usingWhen(
                acquire(),
                resource -> stage.get(),
                this::release,
                this::release);
    }

    private Mono acquire() {
        return Mono
                .fromCallable(() -> {
                    try {
                        outer.acquire();
                        if (delayNanos() > 0) {
                            try {
                                inner.acquire();
                                return new Resource(outer, inner);
                            } catch (InterruptedException e) {
                                throw Exceptions.propagate(e);
                            }
                        }
                        return new Resource(outer, null);
                    } catch (InterruptedException e) {
                        throw Exceptions.propagate(e);
                    }
                })
                .delayUntil(resource -> onComplete());
    }

    private Mono release(Resource resource) {
        return Mono.fromRunnable(() -> {
            if (resource.inner != null) {
                resource.inner.release();
            }
            if (resource.outer != null) {
                resource.outer.release();
            }
        });
    }

    static class Resource {
        private final Semaphore outer;
        private final Semaphore inner;

        Resource(Semaphore outer, Semaphore inner) {
            this.outer = outer;
            this.inner = inner;
        }
    }
}