Redis+Lua脚本生成自增主键

一.简介

从redis的2.6版本开始，增加了对lua脚本的支持，目前可用的库只有如下几个：

1. base lib
2. table lib
3. string lib
4. math lib
5. struct lib
6. cjson lib
7. cmsgpack lib
8. bitop lib
9. redis.sha1hex function

像在redis的lua脚本中执行os.time()会报错，不支持os类库。

二.生成自增主键

1.lua脚本

local temp = 0
local key = KEYS[1]
local inc = ARGV[1]
local exp = ARGV[2]
if inc == nil then
	inc = 1
end
if redis.call('EXISTS',key) == 1 then
	temp = redis.call('GET',key) + inc
	redis.call('SETEX',key,exp,temp)
	return temp
else
	temp = temp + inc
	redis.call('SETEX',key,exp,temp)
	return temp
end

2.通过redis的evalsha调用lua脚本生成自增主键

package com.icheetor.jedis;

import java.util.Arrays;
import java.util.List;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.util.Pool;

public class CreateId {

	private static Pool pool = null;
	private static String create_id_sha = null;
	
	static{
		pool = new JedisPool("127.0.0.1", 6379);
		initCreateIdSha();
	}
	private static void initCreateIdSha() {
		Jedis jedis = pool.getResource();
		create_id_sha = jedis.scriptLoad(LoadLuaScriptUtil.loadLuaScript("/createId.lua"));
		jedis.close();
	}
	/**
	 * 创建自增为1的id
	 * @param sha1
	 * @param key
	 * @return
	 */
	private static String getId(String sha1,String key){
		List keyList = Arrays.asList(key);
		List valueList = Arrays.asList("1","5");
		Jedis jedis = null;
		try {
			jedis = pool.getResource();
			String res = jedis.evalsha(sha1, keyList, valueList).toString();
			return res;
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			if(jedis != null){
				jedis.close();
			}
		}
		return "0";
	}
	public static String getCreateId(String key) {
		if(create_id_sha == null) {
			Jedis jedis = pool.getResource();
			create_id_sha = jedis.scriptLoad(LoadLuaScriptUtil.loadLuaScript("/createId.lua"));
			jedis.close();
		}
		return getId(create_id_sha, key);
	}
	public static void main(String[] args) {
		System.out.println(getCreateId("tbl_test"));
	}
}

evalsha需要提前将lua脚本加载到内存中，生成一串sha，来提高执行效率。

将lua脚本加载到内存中，生成一个大字符串，在通过scriptLoad方法加载生成的大字符串，生成一串sha。加载lua脚本的方法如下，共有两中方式，一种是采用apache的commons-io，另一种通过Google的guava工具类，代码如下：

package com.icheetor.jedis;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.Reader;

import com.google.common.io.CharStreams;

/**
 * 加载lua脚本，生成脚本字符串
 * @author iCheetor
 *
 */
public class LoadLuaScriptUtil {

	/**
	 * 通过apache的commons-io转换为脚本字符串
	 * @param path
	 * @return
	 */
	public static String loadLuaScript(String path){
		String luaStr = null;
		InputStream in = LoadLuaScriptUtil.class.getResourceAsStream(path);
		Reader r = new InputStreamReader(in);
		try {
			luaStr = org.apache.commons.io.IOUtils.toString(r);
		} catch (IOException e) {
			e.printStackTrace();
		}
		return luaStr;
	}
	/**
	 * 通过guava转换为脚本字符串
	 * @param path
	 * @return
	 */
	public static String loadLuaScriptByGuava(String path){
		String luaStr = null;
		InputStream in = LoadLuaScriptUtil.class.getResourceAsStream(path);
		Reader r = new InputStreamReader(in);
		try {
			luaStr = CharStreams.toString(r);
		} catch (IOException e) {
			e.printStackTrace();
		}
		return luaStr;
	}
}

由于redis是单线程的，所以能很好的保证id的自增，在lua脚本中进行的计算，赋值，读取操作也就具有了原子性。