package valkey import ( "context" "crypto/sha1" "encoding/hex" "runtime" "sync/atomic" "github.com/valkey-io/valkey-go/internal/util" ) // NewLuaScript creates a Lua instance whose Lua.Exec uses EVALSHA and EVAL. func NewLuaScript(script string) *Lua { sum := sha1.Sum([]byte(script)) return &Lua{script: script, sha1: hex.EncodeToString(sum[:]), maxp: runtime.GOMAXPROCS(0)} } // NewLuaScriptReadOnly creates a Lua instance whose Lua.Exec uses EVALSHA_RO and EVAL_RO. func NewLuaScriptReadOnly(script string) *Lua { lua := NewLuaScript(script) lua.readonly = true return lua } // Lua represents a valkey lua script. It should be created from the NewLuaScript() or NewLuaScriptReadOnly() type Lua struct { script string sha1 string maxp int readonly bool } // Exec the script to the given Client. // It will first try with the EVALSHA/EVALSHA_RO and then EVAL/EVAL_RO if first try failed. // Cross slot keys are prohibited if the Client is a cluster client. func (s *Lua) Exec(ctx context.Context, c Client, keys, args []string) (resp ValkeyResult) { if s.readonly { resp = c.Do(ctx, c.B().EvalshaRo().Sha1(s.sha1).Numkeys(int64(len(keys))).Key(keys...).Arg(args...).Build()) } else { resp = c.Do(ctx, c.B().Evalsha().Sha1(s.sha1).Numkeys(int64(len(keys))).Key(keys...).Arg(args...).Build()) } if err, ok := IsValkeyErr(resp.Error()); ok && err.IsNoScript() { if s.readonly { resp = c.Do(ctx, c.B().EvalRo().Script(s.script).Numkeys(int64(len(keys))).Key(keys...).Arg(args...).Build()) } else { resp = c.Do(ctx, c.B().Eval().Script(s.script).Numkeys(int64(len(keys))).Key(keys...).Arg(args...).Build()) } } return resp } // LuaExec is a single execution unit of Lua.ExecMulti type LuaExec struct { Keys []string Args []string } // ExecMulti exec the script multiple times by the provided LuaExec to the given Client. // It will first try SCRIPT LOAD the script to all valkey nodes and then exec it with the EVALSHA/EVALSHA_RO. // Cross slot keys within single LuaExec are prohibited if the Client is a cluster client. func (s *Lua) ExecMulti(ctx context.Context, c Client, multi ...LuaExec) (resp []ValkeyResult) { var e atomic.Value util.ParallelVals(s.maxp, c.Nodes(), func(n Client) { if err := n.Do(ctx, n.B().ScriptLoad().Script(s.script).Build()).Error(); err != nil { e.CompareAndSwap(nil, &errs{error: err}) } }) if err := e.Load(); err != nil { resp = make([]ValkeyResult, len(multi)) for i := 0; i < len(resp); i++ { resp[i] = newErrResult(err.(*errs).error) } return } cmds := make(Commands, 0, len(multi)) if s.readonly { for _, m := range multi { cmds = append(cmds, c.B().EvalshaRo().Sha1(s.sha1).Numkeys(int64(len(m.Keys))).Key(m.Keys...).Arg(m.Args...).Build()) } } else { for _, m := range multi { cmds = append(cmds, c.B().Evalsha().Sha1(s.sha1).Numkeys(int64(len(m.Keys))).Key(m.Keys...).Arg(m.Args...).Build()) } } return c.DoMulti(ctx, cmds...) }