com.salesforce.dockerfileimageupdate.utils.RateLimiter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of dockerfile-image-update Show documentation

This tool provides a mechanism to make security updates to docker images at scale. The tool searches github for declared docker images and sends pull requests to projects that are not using the desired version of the requested docker image.

There is a newer version: 1.1.26

Show newest version

package com.salesforce.dockerfileimageupdate.utils;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import io.github.bucket4j.TimeMeter;
import java.util.UnknownFormatConversionException;
import net.sourceforge.argparse4j.inf.Namespace;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * This class uses Token bucket
 * algorithm to evenly process the available resources in a given amount of time.
 *  * With every PR sent by DFIU tool, a token will be consumed. This will ensure that a PR will
 * be raised only when there is a token left in the bucket. Initially bucket will have a fixed
 * number of tokens equal to the 'rateLimit' value. The bucket will be refilled with tokens at
 * the end of every 'rateLimitDuration' with 'rateLimit' tokens ensuring that at any point of
 * time, the number of tokens in the bucket won't exceed the 'rateLimit', limiting the rate at
 * which PRs would be raised. To keep the refill rate uniform, one token will be added to the
 * bucket every 'tokenAddingRate'. If no tokens are left in the bucket, then PR won't be raised
 * and program will halt until next token is available.
 */
public class RateLimiter {

    private static final Logger log = LoggerFactory.getLogger(RateLimiter.class);
    private final Bucket bucket;

    /**
     * constructor to initialize RateLimiter with required values.
     * If time meter is not specified, the system default will be chosen.
     * @param rateLimit RateLimit object
     * @param customTimeMeter custom time metere. should be of type TimeMeter
     * @param  Implementing class for customTimeMeter must be a class
     *           extending TimeMeter
     * @see RateLimit
     * @see TimeMeter
     * @see TimeMeter#SYSTEM_MILLISECONDS
     */
    public  RateLimiter(RateLimit rateLimit, T customTimeMeter) {
        customTimeMeter = customTimeMeter != null ? customTimeMeter : (T) TimeMeter.SYSTEM_MILLISECONDS;
        // refill the bucket at the end of every 'rateLimitDuration' with 'rateLimit' tokens,
        // not exceeding the max capacity
        Refill refill = Refill.intervally(rateLimit.getRate(), rateLimit.getDuration());
        // initially bucket will have no. of tokens equal to its max capacity i.e.
        // the value of 'rateLimit'
        Bandwidth limit = Bandwidth.classic(rateLimit.getRate(), refill);
        this.bucket = Bucket.builder()
                .addLimit(limit)
                // this is internal limit to avoid spikes and distribute the load uniformly over
                // DEFAULT_RATE_LIMIT_DURATION
                // one token added per DEFAULT_TOKEN_ADDING_RATE
                .addLimit(Bandwidth.classic(1, Refill.intervally(1, rateLimit.getTokenAddingRate())))
                .withCustomTimePrecision(customTimeMeter)
                .build();
    }

    public RateLimiter(RateLimit rateLimit) {
        this(rateLimit, TimeMeter.SYSTEM_MILLISECONDS);
    }

    public RateLimiter() {
        this(new RateLimit());
    }

    /**
     * This method will create and return an object of type RateLimiter based on the
     * variable set in the Namespace.
     *
     * @param ns Namespace
     * @return RateLimiter object if required variable is set in Namespace
     * null otherwise.
     * @see net.sourceforge.argparse4j.inf.Namespace Namespace
     */
    public static RateLimiter getInstance(Namespace ns) {
        String rateLimitPRCreation = ns.get(Constants.RATE_LIMIT_PR_CREATION);
        RateLimiter rateLimiter = null;

        if (rateLimitPRCreation != null) {
            try {
                RateLimit rl = RateLimit.tokenizeAndGetRateLimit(rateLimitPRCreation);
                rateLimiter = new RateLimiter(rl);
                log.info("Use rateLimiting is enabled, the PRs will be throttled in this run..");
            } catch (UnknownFormatConversionException ex) {
                log.error("Failed to parse ratelimiting argument, will not impose any rate limits", ex);
            }
        } else {
            log.info("Use rateLimiting is disabled, the PRs will not be throttled in this run..");
        }
        return rateLimiter;
    }

    /**
     * This method consumes a token from the bucked.
     * It blocks the execution when no token is available
     * for consumption.
     *
     * @throws InterruptedException when interrupted
     * @see RateLimiter#bucket
     */
    public void consume() throws InterruptedException {
        bucket.asBlocking().consume(Constants.DEFAULT_CONSUMING_TOKEN_RATE);
    }
}