func StartKV(db *DB, pool *redis.Pool) {
kvDB = db
redisPool = pool
c := redisPool.Get()
defer c.Close()
//cmdSetKV: key, value, expireTime
cmdSetKV = redis.NewScript(3, SCRIPT_SET_KV)
err := cmdSetKV.Load(c)
PanicIfError(err)
//cmdGetDel: key, value, expireTime
cmdGetDel = redis.NewScript(1, SCRIPT_GETDEL)
err = cmdGetDel.Load(c)
PanicIfError(err)
//cmdGetExpiredKV: key, value, expireTime
cmdGetExpiredKV = redis.NewScript(1, SCRIPT_GET_EXPIRED_KV)
err = cmdGetExpiredKV.Load(c)
PanicIfError(err)
//start save to db task
go RepeatSingletonTask(redisPool, "kvSaveToDbTask", saveToDBTask)
//go RepeatSingletonTask(redisPool, "hkvSaveToDB", hkvSaveToDB)
}
// Each lock will have a name which corresponds to a key in the Redis server.
// The mutex will also be initialized with a uuid. The mutex uuid
// can be used to extend the TTL for the lock.
func NewMutex(name string, ttl time.Duration) *Mutex {
m := new(Mutex)
m.Name = name
m.Ttl = ttl
m.Backoff = time.Second
m.id = uuid()
m.lock = redis.NewScript(1, luaLock)
m.unlock = redis.NewScript(1, luaUnlock)
return m
}
func init() {
genericScript = strings.NewReplacer(
"INSERTSUFFIX", insertSuffix,
"DELETESUFFIX", deleteSuffix,
).Replace(genericScript)
insertScript = redis.NewScript(1, strings.NewReplacer(
"REMSUFFIX", deleteSuffix, // Insert script does ZREM from deletes key
"ADDSUFFIX", insertSuffix, // and ZADD to inserts key
).Replace(genericScript))
deleteScript = redis.NewScript(1, strings.NewReplacer(
"REMSUFFIX", insertSuffix, // Delete script does ZREM from inserts key
"ADDSUFFIX", deleteSuffix, // and ZADD to deletes key
).Replace(genericScript))
}
// ReadAllTasks load all unfinished task
func ReadAllTasks() []Task {
c := Pool.Get()
defer c.Close()
var GetAllTasksLua = `
local data = redis.call("LRANGE", "allTasks", "0", "-1")
local ret = {}
for idx=1, #data do
ret[idx] = redis.call("HGETALL", "task:"..data[idx])
end
return ret
`
var tasks []Task
script := redis.NewScript(0, GetAllTasksLua)
values, err := redis.Values(script.Do(c))
if err != nil {
log.Println(err)
}
for i := range values {
t := new(Task)
redis.ScanStruct(values[i].([]interface{}), t)
tasks = append(tasks, *t)
}
return tasks
}
func (rli *RedisInput) Init(config interface{}) error {
rli.conf = config.(*RedisInputConfig)
var err error
rli.conn, err = redis.Dial("tcp", rli.conf.Address)
if err != nil {
return fmt.Errorf("connecting to - %s", err.Error())
}
if _, err = rli.conn.Do("AUTH", rli.conf.Password); err != nil {
rli.conn.Close()
return fmt.Errorf("AUTH - %s", err.Error())
}
rli.batchlpop = redis.NewScript(1, `
local i = tonumber(ARGV[1])
local res = {}
local length = redis.call('llen',KEYS[1])
if length < i then i = length end
while (i > 0) do
local item = redis.call("lpop", KEYS[1])
if (not item) then
break
end
table.insert(res, item)
i = i-1
end
return res
`)
return nil
}
func getFreePort(c *redis.Conn) int {
// refreshOccupiedPorts(c)
s := redis.NewScript(3, `
redis.call("SDIFFSTORE", KEYS[3], KEYS[1], KEYS[2])
redis.call("ZINTERSTORE", KEYS[3], 1, KEYS[3])
local freePort = redis.call("ZRANGE", KEYS[3], 0, 0)[1]
redis.call("SADD", KEYS[2], freePort)
if freePort == nil then
return 0
else
return freePort
end
`)
r, e := redis.Int(s.Do(*c,
getKey("ports", "possible"),
getKey("ports", "occupied"),
getKey("ports", "available")))
if e != nil {
log.Panic("Unable to fetch a free port: ", e)
}
return r
}
func setTags(prefix string, metrics string, tags []string) string {
//TODO: call function
client := clientFunction()
defer client.Close()
hashName := prefix + metrics + "\tTagHash"
listName := prefix + metrics + "\tTagList"
args := []string{}
args = append(args, tags...)
args = append(args, hashName)
args = append(args, listName)
scriptArgs := make([]interface{}, len(args))
for i, v := range args {
scriptArgs[i] = v
}
s := getLuaScript("setTag")
script := redis.NewScript(len(tags), s)
result, err := redis.String(script.Do(client, scriptArgs...))
if err != nil {
log.Println(err)
}
return result
}
func (this __red__model_ServiceType__field_SomeInt) Set(value int) {
conn := this.__.(redis.Conn)
{
value := __red__to_string(value)
_, err := redis.NewScript(0, `
local rut = {}
do
local sort = function (arr) table.sort(arr, function (a, b) return a[1] < b[1] end); return arr end
rut.pack = function (v) return cmsgpack.pack(sort(v)) end
rut.hash = function (v) return redis.sha1hex(rut.pack(v)) end
end
local prim = ARGV[1]
local value = ARGV[2]
local baseKey = "red::obj::ServiceType::base::" .. prim
local set = function ()
redis.call("hset", baseKey, "SomeInt", value)
end
do
--> Unique index
local _SomeInt = redis.call("hget", baseKey, "SomeInt") or ""
local oldKey = "red::idx::ServiceType::SomeInt::" .. rut.hash({
{"SomeInt", _SomeInt},
})
local newKey = "red::idx::ServiceType::SomeInt::" .. rut.hash({
{"SomeInt", value },
})
-- TODO: Ensure uniqueness
-- TODO: What to do if key exists?
if redis.call("exists", oldKey) then
redis.call("set", newKey, prim)
set()
redis.call("del", oldKey)
else
redis.call("set", newKey, prim)
set()
end
end
`).Do(conn, this.__model.__prim, value, __red__to_string(time.Now().UTC().Unix()))
if err != nil {
panic(err)
}
}
}
func InitRedix(addr string) {
Redix = make([]redis.Conn, _Max)
RedixMu = make([]*sync.Mutex, _Max)
var err error
for i := 0; i < _Max; i++ {
Redix[i], err = redis.Dial("tcp", addr)
if err != nil {
panic(err)
}
RedixMu[i] = &sync.Mutex{}
glog.Infof("RedisPool[%d] Init OK on %s\n", i, addr)
}
RedixAddr = addr
netAddr, err := net.ResolveTCPAddr("tcp", addr)
if err != nil {
panic(err)
}
RedixAddrPool = []net.Addr{netAddr}
ScriptSelectMobileId = redis.NewScript(1, _scriptSelectMobileId)
ScriptSelectMobileId.Load(Redix[_SelectMobileId])
ScriptOnline = redis.NewScript(2, _scriptOnline)
ScriptOnline.Load(Redix[_SetUserOnline])
ScriptOffline = redis.NewScript(2, _scriptOffline)
ScriptOffline.Load(Redix[_SetUserOffline])
if gCometType != msgs.CometUdp || gCometUdpSubBindingEvent {
err = SubDeviceUsers()
if err != nil {
panic(err)
}
}
err = SubModifiedPasswd()
if err != nil {
panic(err)
}
if gCometType != msgs.CometUdp || gCometPushUdp {
err = SubCommonMsg()
if err != nil {
panic(err)
}
}
}
func ExampleScript(c redis.Conn, reply interface{}, err error) {
// Initialize a package-level variable with a script.
var getScript = redis.NewScript(1, `return redis.call('get', KEYS[1])`)
// In a function, use the script Do method to evaluate the script. The Do
// method optimistically uses the EVALSHA command. If the script is not
// loaded, then the Do method falls back to the EVAL command.
reply, err = getScript.Do(c, "foo")
}
func TestScript(t *testing.T) {
c, err := redistest.Dial()
if err != nil {
t.Fatalf("error connection to database, %v", err)
}
defer c.Close()
// To test fall back in Do, we make script unique by adding comment with current time.
script := fmt.Sprintf("--%d\nreturn {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}", time.Now().UnixNano())
s := redis.NewScript(2, script)
reply := []interface{}{[]byte("key1"), []byte("key2"), []byte("arg1"), []byte("arg2")}
v, err := s.Do(c, "key1", "key2", "arg1", "arg2")
if err != nil {
t.Errorf("s.Do(c, ...) returned %v", err)
}
if !reflect.DeepEqual(v, reply) {
t.Errorf("s.Do(c, ..); = %v, want %v", v, reply)
}
err = s.Load(c)
if err != nil {
t.Errorf("s.Load(c) returned %v", err)
}
err = s.SendHash(c, "key1", "key2", "arg1", "arg2")
if err != nil {
t.Errorf("s.SendHash(c, ...) returned %v", err)
}
err = c.Flush()
if err != nil {
t.Errorf("c.Flush() returned %v", err)
}
v, err = c.Receive()
if !reflect.DeepEqual(v, reply) {
t.Errorf("s.SendHash(c, ..); c.Receive() = %v, want %v", v, reply)
}
err = s.Send(c, "key1", "key2", "arg1", "arg2")
if err != nil {
t.Errorf("s.Send(c, ...) returned %v", err)
}
err = c.Flush()
if err != nil {
t.Errorf("c.Flush() returned %v", err)
}
v, err = c.Receive()
if !reflect.DeepEqual(v, reply) {
t.Errorf("s.Send(c, ..); c.Receive() = %v, want %v", v, reply)
}
}
func LPOPRPUSH(cnx redis.Conn) *redis.Script {
lmu.Lock()
defer lmu.Unlock()
if lpoprpush == nil {
lpoprpush = redis.NewScript(2, lpoprpushSrc)
}
return lpoprpush
}
func (c RedisProviderConfig) NewProvider() (Provider, error) {
p := &RedisProvider{
BaseProvider: BaseProvider{c},
}
p.pool = redis.Pool{
MaxIdle: 10,
IdleTimeout: 240 * time.Second,
Dial: func() (redis.Conn, error) {
conn, err := redis.Dial("tcp", c.Host+":"+strconv.Itoa(c.Port))
if err != nil {
return nil, err
}
if len(c.Password) > 0 {
if _, err := conn.Do("AUTH", c.Password); err != nil {
conn.Close()
return nil, err
}
}
if _, err := conn.Do("SELECT", c.Database); err != nil {
conn.Close()
return nil, err
}
return conn, err
},
TestOnBorrow: func(c redis.Conn, t time.Time) error {
_, err := c.Do("PING")
return err
},
}
p.countScript = redis.NewScript(0, `return #redis.call('keys', '::till:value:*')`)
p.randomValueKeyScript = redis.NewScript(0, `
local keys = redis.call('keys', '::till:value:*')
if #keys > 0 then
return keys[1]
else
return nil
end
`)
return p, nil
}
// CreateMemo saves a memo
func CreateMemo(user, when, memo string) {
c := Pool.Get()
defer c.Close()
var setMemoLua = `
local id = redis.call("INCR", "memoIncrId")
redis.call("RPUSH", KEYS[1]..":memos", id)
redis.call("HMSET", "memo:"..id, "time", KEYS[2], "content", KEYS[3])
`
script := redis.NewScript(3, setMemoLua)
script.Do(c, user, when, memo)
}
// DeleteMemo deletes a memo
func DeleteMemo(user string, index int) {
c := Pool.Get()
defer c.Close()
var deleteMemoLua = `
local id = redis.call("LINDEX", KEYS[1]..":memos", KEYS[2])
redis.call("LREM", KEYS[1]..":memos", 1, id)
redis.call("DEL", "memo:"..id)
`
script := redis.NewScript(2, deleteMemoLua)
script.Do(c, user, index)
}
// CreateTask saves a task
func CreateTask(ts Task) {
c := Pool.Get()
defer c.Close()
log.Println("save task")
var setTaskLua = `
redis.call("RPUSH", "allTasks", KEYS[1])
redis.call("RPUSH", KEYS[2]..":tasks", KEYS[1])
redis.call("HMSET", "task:"..KEYS[1], "id", KEYS[1], "owner", KEYS[2], "time", KEYS[3], "content", KEYS[4], "chatID", KEYS[5])
`
script := redis.NewScript(5, setTaskLua)
script.Do(c, ts.Id, ts.Owner, ts.When, ts.Desc, ts.ChatId)
}
func ShowOrdersController(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
r.ParseForm()
var accessToken string
if len(r.Form["access_token"]) > 0 {
accessToken = r.Form["access_token"][0]
} else {
accessToken = r.Header.Get("Access-Token")
}
if len(accessToken) <= 0 {
w.WriteHeader(401)
fmt.Fprintf(w, `{"code": "INVALID_ACCESS_TOKEN", "message": "无效的令牌"}`)
return
}
lua := `local uid = redis.call("get", "u:"..KEYS[1])
local oid = redis.call("get", "u:o:"..uid)
if oid then
local menu = redis.call("hgetall", "o:f:"..oid)
return {uid, oid, menu}
end`
var getScript = redis.NewScript(1, lua)
c := rp.Get()
reply, _ := redis.Values(getScript.Do(c, accessToken))
var json bytes.Buffer
if len(reply) != 0 {
var userId, orderId string
var items []interface{}
redis.Scan(reply, &userId, &orderId, &items)
var total int = 0
var j []string
v, _ := redis.Ints(items, nil)
for i := 0; i < len(v); i += 2 {
fid := v[i]
cnt := v[i+1]
total += gfoods[fid].Price * cnt
j = append(j, fmt.Sprintf(`{"food_id": %d, "count": %d}`, fid, cnt))
}
json.WriteString("[{")
json.WriteString(fmt.Sprintf(`"id": "%s", `, orderId))
json.WriteString(fmt.Sprintf(`"items": [%s]`, strings.Join(j, ",")))
json.WriteString(fmt.Sprintf(`,"total": %d`, total))
json.WriteString("}]")
} else {
json.WriteString("{}")
}
w.WriteHeader(200)
fmt.Fprintf(w, json.String())
}
// DeleteTask deletes a task
func DeleteTask(ts Task) {
c := Pool.Get()
defer c.Close()
log.Println("remove", ts.Id)
var removeTaskLua = `
redis.call("LREM", "allTasks", 1, KEYS[2])
redis.call("LREM", KEYS[1]..":tasks", 1, KEYS[2])
redis.call("DEL", "task:"..KEYS[2])
`
script := redis.NewScript(2, removeTaskLua)
script.Do(c, ts.Owner, ts.Id)
}
func getMetric(prefix string) (string, error) {
client := clientFunction()
defer client.Close()
dbName := prefix
s := getLuaScript("getMetric")
script := redis.NewScript(0, s)
result, err := redis.String(script.Do(client, dbName))
if err != nil {
log.Println(err)
}
return result, err
}
func getTags(prefix string, metrics string, target string, keyName string) string {
client := clientFunction()
defer client.Close()
listName := prefix + metrics + "\tTagList"
s := getLuaScript("getTag")
script := redis.NewScript(0, s)
result, err := redis.String(script.Do(client, listName, target, keyName))
if err != nil {
log.Println(err)
}
return result
}
func getMetricKeys(prefix string, metrics string) ([]string, error) {
client := clientFunction()
defer client.Close()
name := prefix + metrics + "\t"
s := getLuaScript("getMetricKeys")
script := redis.NewScript(0, s)
ret, err := redis.Strings(script.Do(client, name))
if err != nil {
log.Println(err)
return []string{}, err
}
return ret, err
}
func (l *RedisInput) Run(output chan common.MapStr) error {
logp.Debug("redisinput", "Running Redis Input")
var keysScript = redis.NewScript(1, `return redis.call('KEYS', KEYS[1])`)
go func() {
redisURL := fmt.Sprintf("redis://%s:%d/%d", l.Host, l.Port, l.DB)
dialConnectTimeout := redis.DialConnectTimeout(3 * time.Second)
dialReadTimeout := redis.DialReadTimeout(10 * time.Second)
var backOffCount = 0
var backOffDuration time.Duration = 5 * time.Second
for {
logp.Debug("redisinput", "Connecting to: %s", redisURL)
server, err := redis.DialURL(redisURL, dialConnectTimeout, dialReadTimeout)
if err != nil {
logp.Err("couldn't start listening: " + err.Error())
return
}
logp.Debug("redisinput", "Connected to Redis Server")
reply, err := keysScript.Do(server, "*")
if err != nil {
logp.Err("An error occured while executing KEYS command: %s\n", err)
return
}
keys, err := redis.Strings(reply, err)
if err != nil {
logp.Err("An error occured while converting reply to String: %s\n", err)
return
}
for _, key := range keys {
logp.Debug("redisinput", "key is %s", key)
lineCount, err := l.handleConn(server, output, key)
if err == nil {
logp.Debug("redisinput", "Read %v events", lineCount)
backOffCount = 0
backOffDuration = time.Duration(backOffCount) * time.Second
time.Sleep(backOffDuration)
} else {
backOffCount++
backOffDuration = time.Duration(backOffCount) * time.Second
time.Sleep(backOffDuration)
}
}
defer server.Close()
}
}()
return nil
}
func refreshOccupiedPorts(c *redis.Conn) {
s := redis.NewScript(2, `
local tenants = redis.call("SMEMBERS", KEYS[1])
redis.call("DEL", KEYS[2])
for i = 1, #tenants do
local port = redis.call("GET", ARGV[1] .. ":" .. tenants[i] .. ":port")
redis.call("SADD", KEYS[2], port)
end
`)
if _, e := s.Do(*c,
getKey("tenants"),
getKey("ports", "occupied"),
getKey("tenant")); e != nil {
log.Panic("Unable to refresh occupied ports", e)
}
}
func loadScripts() {
scriptNames, err := ioutil.ReadDir("lua")
if err != nil {
log.Fatal("could not load lua scripts: ", err)
}
for _, s := range scriptNames {
scriptName := s.Name()
scriptData, err := ioutil.ReadFile("lua/" + scriptName)
if err != nil {
log.Fatal("could not load lua script: ", err)
}
log.Println("loaded script: ", scriptName)
scripts[scriptName] = redis.NewScript(-1, string(scriptData))
}
}
func TestMsgPack(t *testing.T) {
conn, err := redis.Dial("unix", "/tmp/redis.sock")
ferr(err, t)
_, err = conn.Do("set", "key1", "1")
ferr(err, t)
result, err := redis.NewScript(0, `
local key1 = redis.call("get", "key1")
local keys = {{"hello", "kitty"}, {"you are the", key1}, {ARGV[1], "bear"}}
--redis.call("multi")
--redis.call("incr", "_1")
--redis.call("incr", "_2")
--redis.call("exec")
table.sort(keys, function (a, b) return a[1] < b[1] end)
local msg = cmsgpack.pack(keys)
return {msg, redis.sha1hex(msg)}
`).Do(conn, "i a\n\tm русский")
ferr(err, t)
kk := __red__pairs{{"hello", "kitty"}, {"you are the", "1"}, {"i a\n\tm русский", "bear"}}
msgRedis := []string{
string(result.([]interface{})[0].([]byte)),
string(result.([]interface{})[1].([]byte)),
}
log.Println(msgRedis)
log.Println(string(kk.pack()), kk.sha1hex())
if string(kk.pack()) != msgRedis[0] {
t.Log("msgs is not equal")
t.FailNow()
}
if kk.sha1hex() != msgRedis[1] {
t.Log("shas is not equal")
t.FailNow()
}
}
func (c *Conn) loadScript(assetName string) (scriptName string, script *redis.Script, err error) {
log.Printf("Loading script: %s", assetName)
nameWithArity := strings.TrimSuffix(assetName, ".lua")
i := strings.LastIndex(nameWithArity, "_")
scriptName, arityStr := nameWithArity[:i], nameWithArity[i+1:]
content, err := Asset(assetName)
if err != nil {
return
}
arity, err := strconv.Atoi(arityStr)
if err != nil {
return
}
script = redis.NewScript(arity, string(content))
err = script.Load(c)
return
}
func NewCenterPoolAgent(zkAddr string, redisAddr string) *CenterAgent {
zconn, session, zerr := zookeeper.Dial(zkAddr, 5e9)
if zerr != nil {
return nil
}
event := <-session
if event.State != zookeeper.STATE_CONNECTED {
return nil
}
pool := &redis.Pool{
MaxIdle: 1000,
MaxActive: 5000,
IdleTimeout: 20 * time.Second,
Dial: func() (redis.Conn, error) {
connCount = connCount + 1
fmt.Printf("******** connection " + fmt.Sprintf("%d", connCount) + " ********")
return redis.Dial("tcp", redisAddr)
},
}
s := redis.NewScript(1, COUNT_SCRIPT)
return &CenterAgent{zconn, pool, s}
}
func prepareDb() {
c := dial()
defer (*c).Close()
s := redis.NewScript(1, `
redis.call("DEL", KEYS[1])
local portBase = tonumber(ARGV[1])
local maxPorts = tonumber(ARGV[2])
for i = 0, maxPorts - 1 do
redis.call("SADD", KEYS[1], portBase + i)
end
`)
if _, e := s.Do(*c,
getKey("ports", "possible"),
cfg.TenantPortBase,
cfg.MaxTenants); e != nil {
log.Panicf("Unable to prepare database: %v", e)
}
refreshOccupiedPorts(c)
}
// loads all the given scripts from the path
func (l *Lua) LoadScripts(path string) error {
err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
if strings.HasSuffix(info.Name(), ".lua") {
src, err := ioutil.ReadFile(path)
if err != nil {
return err
}
script := redis.NewScript(-1, string(src))
err = script.Load(l.conn)
if err != nil {
return err
}
l.f[info.Name()[:len(info.Name())-4]] = script
// log.Println("Loaded: ", info.Name()[:len(info.Name())-4])
}
return nil
})
return err
}
// ReadAllMemo get user's all memos
func ReadAllMemos(user string) []Memo {
c := Pool.Get()
defer c.Close()
var lua = `
local data = redis.call("LRANGE", KEYS[1]..":memos", "0", "-1")
local ret = {}
for idx=1, #data do
ret[idx] = redis.call("HGETALL", "memo:"..data[idx])
end
return ret
`
var memos []Memo
script := redis.NewScript(1, lua)
values, err := redis.Values(script.Do(c, user))
if err != nil {
log.Println(err)
}
for i := range values {
m := new(Memo)
redis.ScanStruct(values[i].([]interface{}), m)
memos = append(memos, *m)
}
return memos
}
