diff --git a/bytecode.go b/bytecode.go index f3049cee..9d4ba457 100644 --- a/bytecode.go +++ b/bytecode.go @@ -16,6 +16,19 @@ type Bytecode struct { Constants []Object } +// Clone of the bytecode suitable for modification without affecting the original. +// New Bytecode itself is independent, but all the contents of it are still shared +// with the original. +// The only thing that is not shared with the original is Constants slice, as it might be updated +// by ReplaceBuiltinModule(), which should be safe for clone. +func (b *Bytecode) Clone() *Bytecode { + return &Bytecode{ + FileSet: b.FileSet, + MainFunction: b.MainFunction, + Constants: append([]Object{}, b.Constants...), + } +} + // Encode writes Bytecode data to the writer. func (b *Bytecode) Encode(w io.Writer) error { enc := gob.NewEncoder(w) @@ -62,6 +75,21 @@ func (b *Bytecode) FormatConstants() (output []string) { return } +// ReplaceBuiltinModule replaces a builtin module with a new one. +// This is helpful for concurrent script execution, when builtin module does not support +// concurrency and you need to provide custom module instance for each script clone. +func (b *Bytecode) ReplaceBuiltinModule(name string, attrs map[string]Object) { + for i, c := range b.Constants { + switch c := c.(type) { + case *ImmutableMap: + modName := inferModuleName(c) + if modName == name { + b.Constants[i] = (&BuiltinModule{Attrs: attrs}).AsImmutableMap(name) + } + } + } +} + // Decode reads Bytecode data from the reader. func (b *Bytecode) Decode(r io.Reader, modules *ModuleMap) error { if modules == nil { diff --git a/compiler_test.go b/compiler_test.go index 5caf2b03..67760e94 100644 --- a/compiler_test.go +++ b/compiler_test.go @@ -1336,6 +1336,89 @@ func TestCompilerSetImportExt_extension_name_validation(t *testing.T) { } } +func TestCompiler_ReplaceBuiltinModule(t *testing.T) { + sharedValues1 := map[string]int{} + sharedValues2 := map[string]int{} + + createSetter := func(vals map[string]int) map[string]tengo.Object { + return map[string]tengo.Object{ + "set": &tengo.UserFunction{ + Value: func(args ...tengo.Object) (tengo.Object, error) { + n, _ := tengo.ToString(args[0]) + i, _ := tengo.ToInt64(args[1]) + vals[n] = int(i) + return nil, nil + }, + }, + } + } + + checkValues := func(vals map[string]int, name string, base int) { + require.Equal(t, base, vals["direct"], fmt.Sprintf("unexpected value in %s['direct']", name)) + require.Equal(t, base+1, vals["srcM1"], fmt.Sprintf("unexpected value in %s['srcM1']", name)) + require.Equal(t, base+2, vals["srcM2"], fmt.Sprintf("unexpected value in %s['srcM2']", name)) + } + + srcM1 := ` + m := import("setter") + export { set: m.set } + ` + + srcM2 := ` + m := import("m1") + export { set: m.set } + ` + + code := ` + ss := import("setter") + m1 := import("m1") + m2 := import("m2") + + ss.set("direct", value) + m1.set("srcM1", value+1) // should update shared value under once again (src module shares builtin module instance) + m2.set("srcM2", value+2) // should again update the same value (nested src modules share the same builtin module instance) + ` + + script := tengo.NewScript([]byte(code)) + + // set values + err := script.Add("value", 0) + require.NoError(t, err, "failed to set value in script") + + modules := stdlib.GetModuleMap() + modules.AddBuiltinModule("setter", createSetter(sharedValues1)) + modules.AddSourceModule("m1", []byte(srcM1)) + modules.AddSourceModule("m2", []byte(srcM2)) + script.SetImports(modules) + + compiled, err := script.Compile() + require.NoError(t, err, "failed to compile script") + + // Check original script uses builtin module it got at start + compiled.Set("value", 10) + err = compiled.Run() + + require.NoError(t, err, "failed to run original compiled script") + checkValues(sharedValues1, "sharedValues1", 10) + + // Check cloned script uses new builtin module after replacement + clone := compiled.Clone() + clone.Set("value", 20) + clone.ReplaceBuiltinModule("setter", createSetter(sharedValues2)) + err = clone.Run() + + require.NoError(t, err, "failed to run cloned compiled script") + checkValues(sharedValues2, "sharedValues2", 20) + + // Check original script still uses old builtin module + compiled.Set("value", 30) + err = compiled.Run() + + require.NoError(t, err, "failed to run original compiled script") + checkValues(sharedValues1, "sharedValues1", 30) + checkValues(sharedValues2, "sharedValues2", 20) +} + func concatInsts(instructions ...[]byte) []byte { var concat []byte for _, i := range instructions { diff --git a/example_test.go b/example_test.go index 5f074d87..b7358f6e 100644 --- a/example_test.go +++ b/example_test.go @@ -44,3 +44,19 @@ each([a, b, c, d], func(x) { // Output: // 22 288 } + +type TestModule struct { + value int +} + +func (t *TestModule) module() *tengo.BuiltinModule { + return &tengo.BuiltinModule{ + Attrs: map[string]tengo.Object{ + "set": &tengo.UserFunction{Value: func(args ...tengo.Object) (tengo.Object, error) { + i, _ := tengo.ToInt64(args[0]) + t.value = int(i) + return nil, nil + }}, + }, + } +} diff --git a/parser/source_file.go b/parser/source_file.go index e9f4b0f5..2f4eea7e 100644 --- a/parser/source_file.go +++ b/parser/source_file.go @@ -26,7 +26,6 @@ func (p SourceFilePos) IsValid() bool { // line valid position without file name and no column (column == 0) // file invalid position with file name // - invalid position without file name -// func (p SourceFilePos) String() string { s := p.Filename if p.IsValid() { diff --git a/script.go b/script.go index 16f2944b..5cff62bb 100644 --- a/script.go +++ b/script.go @@ -138,6 +138,7 @@ func (s *Script) Compile() (*Compiled, error) { bytecode: bytecode, globals: globals, maxAllocs: s.maxAllocs, + fullClone: true, // we do not share bytecode or global indexes with other clones }, nil } @@ -200,6 +201,7 @@ type Compiled struct { globals []Object maxAllocs int64 lock sync.RWMutex + fullClone bool } // Run executes the compiled script in the virtual machine. @@ -255,6 +257,7 @@ func (c *Compiled) Clone() *Compiled { bytecode: c.bytecode, globals: make([]Object, len(c.globals)), maxAllocs: c.maxAllocs, + fullClone: false, // this clone shares bytecode and global indexes with the 'original' } // copy global objects for idx, g := range c.globals { @@ -265,6 +268,19 @@ func (c *Compiled) Clone() *Compiled { return clone } +// ReplaceBuiltinModule replaces a builtin module with a new one. +// This is helpful for concurrent script execution, when builtin module does not support +// concurrency and you need to provide custom module instance for each script clone. +// +// Remember to call .Clone() to get an instance of the script safe for concurrent use. +func (c *Compiled) ReplaceBuiltinModule(name string, attrs map[string]Object) { + c.lock.Lock() + defer c.lock.Unlock() + + c.prepareForModulesUpdate() + c.bytecode.ReplaceBuiltinModule(name, attrs) +} + // IsDefined returns true if the variable name is defined (has value) before or // after the execution. func (c *Compiled) IsDefined(name string) bool { @@ -336,3 +352,24 @@ func (c *Compiled) Set(name string, value interface{}) error { c.globals[idx] = obj return nil } + +func (c *Compiled) prepareForModulesUpdate() { + if c.fullClone { + return + } + + // To safely modify compiled script internals we need to be sure noone shares + // the same bytecode or global indexes. + // We do not do full copy during Clone() call to skip additional memory allocations + // when they are not needed, leaving Clone() call as optional as it was + // before the 'ReplaceBuiltinModule' feature was added. + + indexes := make(map[string]int, len(c.globalIndexes)) + for name, idx := range c.globalIndexes { + indexes[name] = idx + } + c.globalIndexes = indexes + c.bytecode = c.bytecode.Clone() + + c.fullClone = true +}