redis-限流实现

redis-限流实现

原理

利用redis键过期和nx(未存在才设置成功)的特性，注意，不能将ex和nx作为两个命令分开执行!!,目的和分布式锁一样

单机版

LimitRate:

public class LimitRate {

	private static final String HOST = "192.168.3.xx";
	private static final int PORT = 6380;

	private Jedis jedis;
	private String key;
	private int count;
	private int rate;

	public LimitRate(String key, int count, int second) {
		jedis = new Jedis(HOST, PORT);
		this.key = key;
		this.count = count;
		this.rate = second;
	}


	public boolean acquire() {
		try {
			String ky = "limit:num" + key; 
            //如果key存在，success返回为null,不然返回为"OK" !!
			String success = jedis.set(ky, "1", "nx", "ex", rate);

			if (success != null || jedis.incr(ky) <= count) {
				return Boolean.TRUE;
			}
		} catch (Exception e) {
			e.printStackTrace();
		}

		return Boolean.FALSE;
	}


}

测试:

@Test
public void limitRateTest() {
    LimitRatePool limitRate = new LimitRatePool("user1", 5, 4);
    limitRate.acquire();
}


@Test
public void rateLimitTest() throws InterruptedException {

    LimitRate limitRate = new LimitRate("user1", 5, 2);
    CountDownLatch latch = new CountDownLatch(1);
    for (int i = 0; i < 20; i++) {
        new Thread(new Runnable() {
            @Override
            public void run() {

                try {
                    latch.await();
                    String name = Thread.currentThread().getName();
                    if (limitRate.acquire()) {

                        TimeUnit.MILLISECONDS.sleep(new Random().nextInt(200));
                        System.out.println(name + " 正常执行...");
                    } else {
                        System.out.println(name + " 限流...");
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }

            }
        }).start();
    }
    latch.countDown();

    TimeUnit.SECONDS.sleep(5);
}

结果,测试用例1正常执行，但是在多线程环境下:

Thread-16 正常执行...

Thread-6 限流...
redis.clients.jedis.exceptions.JedisConnectionException: java.net.SocketException: Socket closed
	at redis.clients.jedis.Connection.connect(Connection.java:207)
	at redis.clients.jedis.BinaryClient.connect(BinaryClient.java:93)
	at redis.clients.jedis.Connection.sendCommand(Connection.java:126)
	at redis.clients.jedis.BinaryClient.set(BinaryClient.java:1115)
	at redis.clients.jedis.Client.set(Client.java:847)
	at redis.clients.jedis.Jedis.set(Jedis.java:3035)
	at com.pinlor.cloud.limit.LimitRate.acquire(LimitRate.java:29)
	at com.pinlor.cloud.limit.LimitRateTest$1.run(LimitRateTest.java:25)

好吧，还报错了，而且这种错误不止一个，检查了一下，逻辑没有问题，那可能是多线程环境下没能及时获取jedis资源的问题了，那就要使用JedisPool了

池化版

为了解决在多线程环境下不能及时获取 Jedis 的问题，升级一下 LimitRate

JedisPoolClient,从 JedisPool 池中获取 Jedis:

public class JedisPoolClient {

	private static final String HOST = "192.168.3.xx";
	private static final int PORT = 6380;
	private static JedisPool jedisPool = null;

	static {

		if (jedisPool == null) {
			JedisPoolConfig jedisPoolConfig = new JedisPoolConfig();
			jedisPoolConfig.setMaxTotal(30);
			jedisPoolConfig.setMaxIdle(5);
			jedisPoolConfig.setMaxWaitMillis(100);
			jedisPoolConfig.setTestOnBorrow(false);
			jedisPoolConfig.setTestOnReturn(true);
			jedisPool = new JedisPool(jedisPoolConfig, HOST, PORT);
		}
	}


	public static Jedis getResource() {
		return jedisPool.getResource();
	}


	public static void returnResource(Jedis jedis) {
		if (jedis != null) {
			jedis.close();
		}
	}

}

限流类 LimitRatePool:

public class LimitRatePool {

	private String key;
	private int count;
	private int rate;

	public LimitRatePool(String key, int count, int second) {
		this.key = key;
		this.count = count;
		this.rate = second;
	}


	public boolean acquire() {
		Jedis jedis = null;
		try {
			String ky = "limit:num" + key;
			jedis = JedisPoolClient.getResource();
			String success = jedis.set(ky, "1", "nx", "ex", rate);
			if (success != null || jedis.incr(ky) <= count) {
				return Boolean.TRUE;
			}
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			JedisPoolClient.returnResource(jedis);
		}

		return Boolean.FALSE;
	}

}

结果:

Thread-2 限流...
Thread-3 限流...
Thread-11 限流...
Thread-18 限流...
Thread-9 限流...
Thread-15 限流...
Thread-12 限流...
Thread-5 限流...
Thread-0 限流...
Thread-4 限流...
Thread-10 限流...
Thread-19 限流...
Thread-14 限流...
Thread-1 限流...
Thread-7 限流...
Thread-13 正常执行...
Thread-6 正常执行...
Thread-17 正常执行...
Thread-8 正常执行...
Thread-16 正常执行...

没有之前的问题，也没有报错，perfect !!