package server import ( "encoding/hex" "fmt" "github.com/aarzilli/golua/lua" "github.com/siddontang/ledisdb/ledis" "io" "sync" ) //ledis <-> lua type conversion, same as http://redis.io/commands/eval type luaWriter struct { l *lua.State } func (w *luaWriter) writeError(err error) { panic(err) } func (w *luaWriter) writeStatus(status string) { w.l.NewTable() top := w.l.GetTop() w.l.PushString("ok") w.l.PushString(status) w.l.SetTable(top) } func (w *luaWriter) writeInteger(n int64) { w.l.PushInteger(n) } func (w *luaWriter) writeBulk(b []byte) { if b == nil { w.l.PushBoolean(false) } else { w.l.PushString(ledis.String(b)) } } func (w *luaWriter) writeArray(lst []interface{}) { if lst == nil { w.l.PushBoolean(false) return } w.l.CreateTable(len(lst), 0) top := w.l.GetTop() for i, _ := range lst { w.l.PushInteger(int64(i) + 1) switch v := lst[i].(type) { case []interface{}: w.writeArray(v) case [][]byte: w.writeSliceArray(v) case []byte: w.writeBulk(v) case nil: w.writeBulk(nil) case int64: w.writeInteger(v) default: panic("invalid array type") } w.l.SetTable(top) } } func (w *luaWriter) writeSliceArray(lst [][]byte) { if lst == nil { w.l.PushBoolean(false) return } w.l.CreateTable(len(lst), 0) for i, v := range lst { w.l.PushString(ledis.String(v)) w.l.RawSeti(-2, i+1) } } func (w *luaWriter) writeFVPairArray(lst []ledis.FVPair) { if lst == nil { w.l.PushBoolean(false) return } w.l.CreateTable(len(lst)*2, 0) for i, v := range lst { w.l.PushString(ledis.String(v.Field)) w.l.RawSeti(-2, 2*i+1) w.l.PushString(ledis.String(v.Value)) w.l.RawSeti(-2, 2*i+2) } } func (w *luaWriter) writeScorePairArray(lst []ledis.ScorePair, withScores bool) { if lst == nil { w.l.PushBoolean(false) return } if withScores { w.l.CreateTable(len(lst)*2, 0) for i, v := range lst { w.l.PushString(ledis.String(v.Member)) w.l.RawSeti(-2, 2*i+1) w.l.PushString(ledis.String(ledis.StrPutInt64(v.Score))) w.l.RawSeti(-2, 2*i+2) } } else { w.l.CreateTable(len(lst), 0) for i, v := range lst { w.l.PushString(ledis.String(v.Member)) w.l.RawSeti(-2, i+1) } } } func (w *luaWriter) writeBulkFrom(n int64, rb io.Reader) { w.writeError(fmt.Errorf("unsupport")) } func (w *luaWriter) flush() { } type script struct { sync.Mutex app *App l *lua.State c *client chunks map[string]struct{} } func (app *App) openScript() { s := new(script) s.app = app s.chunks = make(map[string]struct{}) app.s = s l := lua.NewState() l.OpenBase() l.OpenLibs() l.OpenMath() l.OpenString() l.OpenTable() l.OpenPackage() s.l = l s.c = newClient(app) s.c.db = nil w := new(luaWriter) w.l = l s.c.resp = w l.NewTable() l.PushString("call") l.PushGoFunction(luaCall) l.SetTable(-3) l.PushString("pcall") l.PushGoFunction(luaPCall) l.SetTable(-3) l.PushString("sha1hex") l.PushGoFunction(luaSha1Hex) l.SetTable(-3) l.PushString("error_reply") l.PushGoFunction(luaErrorReply) l.SetTable(-3) l.PushString("status_reply") l.PushGoFunction(luaStatusReply) l.SetTable(-3) l.SetGlobal("ledis") setMapState(l, s) } func (app *App) closeScript() { app.s.l.Close() delMapState(app.s.l) app.s = nil } var mapState = map[*lua.State]*script{} var stateLock sync.Mutex func setMapState(l *lua.State, s *script) { stateLock.Lock() defer stateLock.Unlock() mapState[l] = s } func getMapState(l *lua.State) *script { stateLock.Lock() defer stateLock.Unlock() return mapState[l] } func delMapState(l *lua.State) { stateLock.Lock() defer stateLock.Unlock() delete(mapState, l) } func luaCall(l *lua.State) int { return luaCallGenericCommand(l) } func luaPCall(l *lua.State) (n int) { defer func() { if e := recover(); e != nil { luaPushError(l, fmt.Sprintf("%v", e)) n = 1 } return }() return luaCallGenericCommand(l) } func luaErrorReply(l *lua.State) int { return luaReturnSingleFieldTable(l, "err") } func luaStatusReply(l *lua.State) int { return luaReturnSingleFieldTable(l, "ok") } func luaReturnSingleFieldTable(l *lua.State, filed string) int { if l.GetTop() != 1 || l.Type(-1) != lua.LUA_TSTRING { luaPushError(l, "wrong number or type of arguments") return 1 } l.NewTable() l.PushString(filed) l.PushValue(-3) l.SetTable(-3) return 1 } func luaSha1Hex(l *lua.State) int { argc := l.GetTop() if argc != 1 { luaPushError(l, "wrong number of arguments") return 1 } s := l.ToString(1) s = hex.EncodeToString(ledis.Slice(s)) l.PushString(s) return 1 } func luaPushError(l *lua.State, msg string) { l.NewTable() l.PushString("err") err := l.NewError(msg) l.PushString(err.Error()) l.SetTable(-3) } func luaCallGenericCommand(l *lua.State) int { s := getMapState(l) if s == nil { panic("Invalid lua call") } else if s.c.db == nil { panic("Invalid lua call, not prepared") } c := s.c argc := l.GetTop() if argc < 1 { panic("Please specify at least one argument for ledis.call()") } c.cmd = l.ToString(1) c.args = make([][]byte, argc-1) for i := 2; i <= argc; i++ { switch l.Type(i) { case lua.LUA_TNUMBER: c.args[i-2] = []byte(fmt.Sprintf("%.17g", l.ToNumber(i))) case lua.LUA_TSTRING: c.args[i-2] = []byte(l.ToString(i)) default: panic("Lua ledis() command arguments must be strings or integers") } } c.perform() return 1 } func luaSetGlobalArray(l *lua.State, name string, ay [][]byte) { l.NewTable() for i := 0; i < len(ay); i++ { l.PushString(ledis.String(ay[i])) l.RawSeti(-2, i+1) } l.SetGlobal(name) } func luaReplyToLedisReply(l *lua.State) interface{} { base := l.GetTop() defer func() { l.SetTop(base - 1) }() switch l.Type(-1) { case lua.LUA_TSTRING: return ledis.Slice(l.ToString(-1)) case lua.LUA_TBOOLEAN: if l.ToBoolean(-1) { return int64(1) } else { return nil } case lua.LUA_TNUMBER: return int64(l.ToInteger(-1)) case lua.LUA_TTABLE: l.PushString("err") l.GetTable(-2) if l.Type(-1) == lua.LUA_TSTRING { return fmt.Errorf("%s", l.ToString(-1)) } l.Pop(1) l.PushString("ok") l.GetTable(-2) if l.Type(-1) == lua.LUA_TSTRING { return l.ToString(-1) } else { l.Pop(1) ay := make([]interface{}, 0) for i := 1; ; i++ { l.PushInteger(int64(i)) l.GetTable(-2) if l.Type(-1) == lua.LUA_TNIL { l.Pop(1) break } ay = append(ay, luaReplyToLedisReply(l)) } return ay } default: return nil } }