TokenBuckets.java 3.94 KB
package com.zhaoonline.coupen.concurrent.tokenbucket;


import com.google.common.base.Ticker;
import com.google.common.util.concurrent.Uninterruptibles;

import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
/**
 * Created by ZhaoOnline<br/>
 * User: yangyoupeng<br/>
 * Date: 2016/12/16<br/>
 * Time: 17:42<br/>
 * Description:please descript you class
 */

/**
 * Static utility methods pertaining to creating {@link TokenBucketImpl} instances.
 */
public final class TokenBuckets {
    private TokenBuckets() {
    }

    /**
     * Create a new builder for token buckets.
     */
    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private Long capacity = null;
        private long initialTokens = 0;
        private TokenBucketImpl.RefillStrategy refillStrategy = null;
        private TokenBucketImpl.SleepStrategy sleepStrategy = YIELDING_SLEEP_STRATEGY;
        private final Ticker ticker = Ticker.systemTicker();

        /**
         * Specify the overall capacity of the token bucket.
         */
        public Builder withCapacity(long numTokens) {
            checkArgument(numTokens > 0, "Must specify a positive number of tokens");
            capacity = numTokens;
            return this;
        }

        /**
         * Initialize the token bucket with a specific number of tokens.
         */
        public Builder withInitialTokens(long numTokens) {
            checkArgument(numTokens > 0, "Must specify a positive number of tokens");
            initialTokens = numTokens;
            return this;
        }

        /**
         * Refill tokens at a fixed interval.
         */
        public Builder withFixedIntervalRefillStrategy(long refillTokens, long period, TimeUnit unit) {
            return withRefillStrategy(new FixedIntervalRefillStrategy(ticker, refillTokens, period, unit));
        }

        /**
         * Use a user defined refill strategy.
         */
        public Builder withRefillStrategy(TokenBucket.RefillStrategy refillStrategy) {
            this.refillStrategy = checkNotNull(refillStrategy);
            return this;
        }

        /**
         * Use a sleep strategy that will always attempt to yield the CPU to other processes.
         */
        public Builder withYieldingSleepStrategy() {
            return withSleepStrategy(YIELDING_SLEEP_STRATEGY);
        }

        /**
         * Use a sleep strategy that will not yield the CPU to other processes.  It will busy wait until more tokens become
         * available.
         */
        public Builder withBusyWaitSleepStrategy() {
            return withSleepStrategy(BUSY_WAIT_SLEEP_STRATEGY);
        }

        /**
         * Use a user defined sleep strategy.
         */
        public Builder withSleepStrategy(TokenBucket.SleepStrategy sleepStrategy) {
            this.sleepStrategy = checkNotNull(sleepStrategy);
            return this;
        }

        /**
         * Build the token bucket.
         */
        public TokenBucket build() {
            checkNotNull(capacity, "Must specify a capacity");
            checkNotNull(refillStrategy, "Must specify a refill strategy");

            return new TokenBucketImpl(capacity, initialTokens, refillStrategy, sleepStrategy);
        }
    }

    private static final TokenBucketImpl.SleepStrategy YIELDING_SLEEP_STRATEGY = new TokenBucketImpl.SleepStrategy() {
        @Override
        public void sleep() {
            // Sleep for the smallest unit of time possible just to relinquish control
            // and to allow other threads to run.
            Uninterruptibles.sleepUninterruptibly(1, TimeUnit.NANOSECONDS);
        }
    };

    private static final TokenBucketImpl.SleepStrategy BUSY_WAIT_SLEEP_STRATEGY = new TokenBucketImpl.SleepStrategy() {
        @Override
        public void sleep() {
            // Do nothing, don't sleep.
        }
    };
}