uw.task.util.GlobalRateLimiter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of uw-task Show documentation
uw-task包是一个分布式任务框架，通过uw-task可以快速构建基于docker的分布式任务体系，同时支持基于influx的任务运维监控。
The newest version!
package uw.task.util;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.support.atomic.RedisAtomicLong;
import org.springframework.stereotype.Component;

/**
 * 基于redis实现的序列管理器。
 * 
 * @author axeon
 *
 */
@Component
public class GlobalRateLimiter {

	private static final Logger log = LoggerFactory.getLogger(GlobalRateLimiter.class);

	private static final String REDIS_TAG = "_RATE_LIMIT_";

	@Autowired
	private StringRedisTemplate template;

	private ConcurrentHashMap map = new ConcurrentHashMap<>();

	/**
	 * 尝试获得限制允许状态。
	 * 
	 * @param name
	 * @return
	 */
	public long tryAcquire(String name) {
		long wait = 0;
		try {
			RedisRateLimiter limiter = map.get(name);
			if (limiter != null) {
				wait = limiter.tryAcquire(1);
			}
		} catch (Exception e) {
			log.error(e.getMessage(), e);
		}
		return wait;
	}

	/**
	 * 初始化一个流量限速器。
	 * 
	 * @param name
	 *            限速器名称，应是全局唯一的名称
	 * @param limitRate
	 *            限速速率
	 * @param limitTimeUnit
	 *            限速时间单位
	 * @param limitTimeValue
	 *            限速时间数值
	 * @return 如果已经存在，则返回false，不再创建。
	 */
	public boolean initLimiter(String name, long limitRate, TimeUnit limitTimeUnit, long limitTimeValue) {
		boolean flag = false;
		try {
			RedisRateLimiter limiter = map.get(name);
			if (limiter == null) {
				limiter = new RedisRateLimiter(name, limitRate, limitTimeUnit, limitTimeValue);
				map.put(name, limiter);
				flag = true;
			} else {
				limiter.setLimitRate(limitRate);
				limiter.setLimitTimeUnit(limitTimeUnit);
				limiter.setLimitTimeValue(limitTimeValue);
				flag = false;
			}
		} catch (Exception e) {
			log.error(e.getMessage(), e);
		}
		return flag;
	}

	/**
	 * redis流量限速器。
	 * 
	 * @author axeon
	 *
	 */
	class RedisRateLimiter {

		/**
		 * 限速器名称
		 */
		String name;
		/**
		 * 流量限速
		 */
		long limitRate = 1l;

		/**
		 * 限速时间单位
		 */
		TimeUnit limitTimeUnit = TimeUnit.SECONDS;

		/**
		 * 限速时间数值
		 */
		long limitTimeValue = 1;

		/**
		 * redis计数器
		 */
		private RedisAtomicLong counter;

		/**
		 * @return the limitRate
		 */
		public long getLimitRate() {
			return limitRate;
		}

		/**
		 * @param limitRate
		 *            the limitRate to set
		 */
		public void setLimitRate(long limitRate) {
			this.limitRate = limitRate;
		}

		/**
		 * @return the limitTimeUnit
		 */
		public TimeUnit getLimitTimeUnit() {
			return limitTimeUnit;
		}

		/**
		 * @param limitTimeUnit
		 *            the limitTimeUnit to set
		 */
		public void setLimitTimeUnit(TimeUnit limitTimeUnit) {
			this.limitTimeUnit = limitTimeUnit;
		}

		/**
		 * @return the limitTimeValue
		 */
		public long getLimitTimeValue() {
			return limitTimeValue;
		}

		/**
		 * @param limitTimeValue
		 *            the limitTimeValue to set
		 */
		public void setLimitTimeValue(long limitTimeValue) {
			this.limitTimeValue = limitTimeValue;
		}

		/**
		 * 初始化一个流量限速器。
		 * 
		 * @param limitRate
		 *            限速速率
		 * @param limitTimeUnit
		 *            限速时间单位
		 * @param limitTimeValue
		 *            限速时间数值
		 */
		public RedisRateLimiter(String name, long limitRate, TimeUnit limitTimeUnit, long limitTimeValue) {
			super();
			this.name = name;
			this.limitRate = limitRate;
			this.limitTimeUnit = limitTimeUnit;
			this.limitTimeValue = limitTimeValue;
			counter = new RedisAtomicLong(REDIS_TAG + name, template.getConnectionFactory());
		}

		/**
		 * 检查是否超限。
		 * 
		 * @param permits
		 * @return 如果为超限则返回0，否则返回需要等待的秒数
		 */
		long tryAcquire(int permits) {
			long expire = -1;
			long value = counter.addAndGet(permits);
			if (value == 1) {
				counter.expire(limitTimeValue, limitTimeUnit);
			}

			if (value > limitRate) {
				expire = counter.getExpire();
				// 防止 value == 1时的expire操作没有成功
				if (expire == -1) {
					counter.expire(limitTimeValue, limitTimeUnit);
				}
			}
			if (value <= limitRate) {
				return 0;
			} else {
				return expire;
			}
		}

	}

}