Allow builtin modules replacement in compiled scripts

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 {

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 <setter> 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 {

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
+			}},
+		},
+	}
+}

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() {

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
+}