added Context.Close() and ModuleImpl.OnContextClosed()

go-python · drew-512 · Dec 8, 2021 · Dec 8, 2021 · Dec 8, 2021 · Dec 8, 2021
commit 681abf2de628248df3cbba21d4e29b80e21e2232
diff --git a/modules/runtime.go b/modules/runtime.go
@@ -11,6 +11,7 @@ import (
 	"path"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"github.com/go-python/gpython/marshal"
 	"github.com/go-python/gpython/py"
@@ -29,14 +30,22 @@ func init() {
 
 // context implements py.Context
 type context struct {
-	store *py.ModuleStore
-	opts  py.ContextOpts
+	store     *py.ModuleStore
+	opts      py.ContextOpts
+	closeOnce sync.Once
+	closing   bool
+	closed    bool
+	running   sync.WaitGroup
+	done      chan struct{}
 }
 
 // See py.Context interface
 func NewContext(opts py.ContextOpts) py.Context {
 	ctx := &context{
-		opts: opts,
+		opts:    opts,
+		done:    make(chan struct{}),
+		closing: false,
+		closed:  false,
 	}
 
 	ctx.store = py.NewModuleStore()
@@ -51,7 +60,11 @@ func NewContext(opts py.ContextOpts) py.Context {
 }
 
 func (ctx *context) ModuleInit(impl *py.ModuleImpl) (*py.Module, error) {
-	var err error
+	err := ctx.PushBusy()
+	defer ctx.PopBusy()
+	if err != nil {
+		return nil, err
+	}
 
 	if impl.Code == nil && len(impl.CodeSrc) > 0 {
 		impl.Code, err = py.Compile(string(impl.CodeSrc), impl.Info.FileDesc, py.ExecMode, 0, true)
@@ -72,7 +85,7 @@ func (ctx *context) ModuleInit(impl *py.ModuleImpl) (*py.Module, error) {
 		}
 	}
 
-	module, err := ctx.Store().NewModule(ctx, impl.Info, impl.Methods, impl.Globals)
+	module, err := ctx.Store().NewModule(ctx, impl)
 	if err != nil {
 		return nil, err
 	}
@@ -88,14 +101,20 @@ func (ctx *context) ModuleInit(impl *py.ModuleImpl) (*py.Module, error) {
 }
 
 func (ctx *context) ResolveAndCompile(pathname string, opts py.CompileOpts) (py.CompileOut, error) {
+	err := ctx.PushBusy()
+	defer ctx.PopBusy()
+	if err != nil {
+		return py.CompileOut{}, err
+	}
+
 	tryPaths := defaultPaths
 	if opts.UseSysPaths {
 		tryPaths = ctx.Store().MustGetModule("sys").Globals["path"].(*py.List).Items
 	}
 
 	out := py.CompileOut{}
 
-	err := resolveRunPath(pathname, opts, tryPaths, func(fpath string) (bool, error) {
+	err = resolveRunPath(pathname, opts, tryPaths, func(fpath string) (bool, error) {
 
 		stat, err := os.Stat(fpath)
 		if err == nil && stat.IsDir() {
@@ -162,6 +181,36 @@ func (ctx *context) ResolveAndCompile(pathname string, opts py.CompileOpts) (py.
 	return out, nil
 }
 
+func (ctx *context) PushBusy() error {
+	if ctx.closed {
+		return py.ExceptionNewf(py.RuntimeError, "Context closed")
+	}
+	ctx.running.Add(1)
+	return nil
+}
+
+func (ctx *context) PopBusy() {
+	ctx.running.Done()
+}
+
+// Close -- see type py.Context
+func (ctx *context) Close() {
+	ctx.closeOnce.Do(func() {
+		ctx.closing = true
+		ctx.running.Wait()
+		ctx.closed = true
+
+		// Give each module a chance to release resources
+		ctx.store.OnContextClosed()
+		close(ctx.done)
+	})
+}
+
+// Done -- see type py.Context
+func (ctx *context) Done() <-chan struct{} {
+	return ctx.done
+}
+
 var defaultPaths = []py.Object{
 	py.String("."),
 }
@@ -216,6 +265,12 @@ func resolveRunPath(runPath string, opts py.CompileOpts, pathObjs []py.Object, t
 }
 
 func (ctx *context) RunCode(code *py.Code, globals, locals py.StringDict, closure py.Tuple) (py.Object, error) {
+	err := ctx.PushBusy()
+	defer ctx.PopBusy()
+	if err != nil {
+		return nil, err
+	}
+
 	return vm.EvalCode(ctx, code, globals, locals, nil, nil, nil, nil, closure)
 }
 

diff --git a/py/module.go b/py/module.go
@@ -14,9 +14,10 @@ import (
 type ModuleFlags int32
 
 const (
-	// Set for modules that are threadsafe, stateless, and/or can be shared across multiple py.Context instances (for efficiency).
-	// Otherwise, a separate module instance is created for each py.Context that imports it.
-	ShareModule ModuleFlags = 0x01 // @@TODO
+	// ShareModule signals that an embedded module is threadsafe and read-only, meaninging it could be shared across multiple py.Context instances (for efficiency).
+	// Otherwise, ModuleImpl will create a separate py.Module instance for each py.Context that imports it.
+	// This should be used with extreme caution since any module mutation (write) means possible cross-context data corruption.
+	ShareModule ModuleFlags = 0x01
 
 	MainModuleName = "__main__"
 )
@@ -34,12 +35,13 @@ type ModuleInfo struct {
 // By convention, .Code is executed when a module instance is initialized.
 // If .Code == nil, then .CodeBuf or .CodeSrc will be auto-compiled to set .Code.
 type ModuleImpl struct {
-	Info    ModuleInfo
-	Methods []*Method
-	Globals StringDict
-	CodeSrc string // Module code body (py source code to be compiled)
-	CodeBuf []byte // Module code body (serialized py.Code object)
-	Code    *Code  // Module code body
+	Info            ModuleInfo
+	Methods         []*Method     // Module-bound global method functions
+	Globals         StringDict    // Module-bound global variables
+	CodeSrc         string        // Module code body (source code to be compiled)
+	CodeBuf         []byte        // Module code body (serialized py.Code object)
+	Code            *Code         // Module code body
+	OnContextClosed func(*Module) // Callback for when a py.Context is closing to release resources
 }
 
 // ModuleStore is a container of Module imported into an owning py.Context.
@@ -87,10 +89,9 @@ func NewModuleStore() *ModuleStore {
 
 // Module is a runtime instance of a ModuleImpl bound to the py.Context that imported it.
 type Module struct {
-	ModuleInfo
-
-	Globals StringDict
-	Context Context
+	ModuleImpl *ModuleImpl // Parent implementation of this Module instance
+	Globals    StringDict  // Initialized from ModuleImpl.Globals
+	Context    Context     // Parent context that "owns" this Module instance
 }
 
 var ModuleType = NewType("module", "module object")
@@ -101,51 +102,68 @@ func (o *Module) Type() *Type {
 }
 
 func (m *Module) M__repr__() (Object, error) {
-	return String(fmt.Sprintf("<module %s>", m.Name)), nil
+	name, ok := m.Globals["__name__"].(String)
+	if !ok {
+		name = "???"
+	}
+	return String(fmt.Sprintf("<module %s>", string(name))), nil
 }
 
 // Get the Dict
 func (m *Module) GetDict() StringDict {
 	return m.Globals
 }
 
+// Calls a named method of a module
+func (m *Module) Call(name string, args Tuple, kwargs StringDict) (Object, error) {
+	attr, err := GetAttrString(m, name)
+	if err != nil {
+		return nil, err
+	}
+	return Call(attr, args, kwargs)
+}
+
+// Interfaces
+var _ IGetDict = (*Module)(nil)
+
 // NewModule adds a new Module instance to this ModuleStore.
 // Each given Method prototype is used to create a new "live" Method bound this the newly created Module.
 // This func also sets appropriate module global attribs based on the given ModuleInfo (e.g. __name__).
-func (store *ModuleStore) NewModule(ctx Context, info ModuleInfo, methods []*Method, globals StringDict) (*Module, error) {
-	if info.Name == "" {
-		info.Name = MainModuleName
+func (store *ModuleStore) NewModule(ctx Context, impl *ModuleImpl) (*Module, error) {
+	name := impl.Info.Name
+	if name == "" {
+		name = MainModuleName
 	}
 	m := &Module{
-		ModuleInfo: info,
-		Globals:    globals.Copy(),
+		ModuleImpl: impl,
+		Globals:    impl.Globals.Copy(),
 		Context:    ctx,
 	}
 	// Insert the methods into the module dictionary
 	// Copy each method an insert each "live" with a ptr back to the module (which can also lead us to the host Context)
-	for _, method := range methods {
+	for _, method := range impl.Methods {
 		methodInst := new(Method)
 		*methodInst = *method
 		methodInst.Module = m
 		m.Globals[method.Name] = methodInst
 	}
 	// Set some module globals
-	m.Globals["__name__"] = String(info.Name)
-	m.Globals["__doc__"] = String(info.Doc)
+	m.Globals["__name__"] = String(name)
+	m.Globals["__doc__"] = String(impl.Info.Doc)
 	m.Globals["__package__"] = None
-	if len(info.FileDesc) > 0 {
-		m.Globals["__file__"] = String(info.FileDesc)
+	if len(impl.Info.FileDesc) > 0 {
+		m.Globals["__file__"] = String(impl.Info.FileDesc)
 	}
 	// Register the module
-	store.modules[info.Name] = m
+	store.modules[name] = m
 	// Make a note of some modules
-	switch info.Name {
+	switch name {
 	case "builtins":
 		store.Builtins = m
 	case "importlib":
 		store.Importlib = m
 	}
-	// fmt.Printf("Registering module %q\n", name)
+	// fmt.Printf("Registered module %q\n", moduleName)
 	return m, nil
 }
 
@@ -167,14 +185,11 @@ func (store *ModuleStore) MustGetModule(name string) *Module {
 	return m
 }
 
-// Calls a named method of a module
-func (m *Module) Call(name string, args Tuple, kwargs StringDict) (Object, error) {
-	attr, err := GetAttrString(m, name)
-	if err != nil {
-		return nil, err
+// OnContextClosed signals all module instances that the parent py.Context has closed
+func (store *ModuleStore) OnContextClosed() {
+	for _, m := range store.modules {
+		if m.ModuleImpl.OnContextClosed != nil {
+			m.ModuleImpl.OnContextClosed(m)
+		}
 	}
-	return Call(attr, args, kwargs)
 }
-
-// Interfaces
-var _ IGetDict = (*Module)(nil)
diff --git a/py/run.go b/py/run.go
@@ -33,15 +33,18 @@ type Context interface {
 	// RunCode is a lower-level invocation to execute the given py.Code.
 	RunCode(code *Code, globals, locals StringDict, closure Tuple) (result Object, err error)
 
-	// Execution of any of the above will stop when the next opcode runs
-	// @@TODO
-	// SignalHalt()
-
 	// Returns the named module for this context (or an error if not found)
 	GetModule(moduleName string) (*Module, error)
 
 	// Gereric access to this context's modules / state.
 	Store() *ModuleStore
+
+	// Close signals that is context is about to go out of scope and any internal resources should be released.
+	// Operations on a py.Context that have closed will generally result in an error.
+	Close()
+
+	// Done returns a signal that can be used to detect when this Context has fully closed / completed.
+	Done() <-chan struct{}
 }
 
 // CompileOpts specifies options for high-level compilation.
@@ -100,20 +103,45 @@ var (
 	Compile func(src, srcDesc string, mode CompileMode, flags int, dont_inherit bool) (*Code, error)
 )
 
-// RunFile resolves the given pathname, compiles as needed, and runs that code in the given module, returning the Module to indicate success.
-// If inModule is a *Module, then the code is run in that module.
-// If inModule is nil, the code is run in a new __main__ module (and the new Module is returned).
-// If inModule is a string, the code is run in a new module with the given name (and the new Module is returned).
+// RunFile resolves the given pathname, compiles as needed, executes the code in the given module, and returns the Module to indicate success.
+//
+// See RunCode() for description of inModule.
 func RunFile(ctx Context, pathname string, opts CompileOpts, inModule interface{}) (*Module, error) {
 	out, err := ctx.ResolveAndCompile(pathname, opts)
 	if err != nil {
 		return nil, err
 	}
 
-	var moduleName string
-	createNew := false
-	var module *Module
+	return RunCode(ctx, out.Code, out.FileDesc, inModule)
+}
+
+// RunSrc compiles the given python buffer and executes it within the given module and returns the Module to indicate success.
+//
+// See RunCode() for description of inModule.
+func RunSrc(ctx Context, pySrc string, pySrcDesc string, inModule interface{}) (*Module, error) {
+	if pySrcDesc == "" {
+		pySrcDesc = "<run>"
+	}
+	code, err := Compile(pySrc+"\n", pySrcDesc, SingleMode, 0, true)
+	if err != nil {
+		return nil, err
+	}
+
+	return RunCode(ctx, code, pySrcDesc, inModule)
+}
 
+// RunCode executes the given code object within the given module and returns the Module to indicate success.
+// If inModule is a *Module, then the code is run in that module.
+// If inModule is nil, the code is run in a new __main__ module (and the new Module is returned).
+// If inModule is a string, the code is run in a new module with the given name (and the new Module is returned).
+func RunCode(ctx Context, code *Code, codeDesc string, inModule interface{}) (*Module, error) {
+	var (
+		module     *Module
+		moduleName string
+		err        error
+	)
+
+	createNew := false
 	switch mod := inModule.(type) {
 
 	case string:
@@ -122,7 +150,7 @@ func RunFile(ctx Context, pathname string, opts CompileOpts, inModule interface{
 	case nil:
 		createNew = true
 	case *Module:
-		_, err = ctx.RunCode(out.Code, mod.Globals, mod.Globals, nil)
+		_, err = ctx.RunCode(code, mod.Globals, mod.Globals, nil)
 		module = mod
 	default:
 		err = ExceptionNewf(TypeError, "unsupported module type: %v", inModule)
@@ -132,9 +160,9 @@ func RunFile(ctx Context, pathname string, opts CompileOpts, inModule interface{
 		moduleImpl := ModuleImpl{
 			Info: ModuleInfo{
 				Name:     moduleName,
-				FileDesc: out.FileDesc,
+				FileDesc: codeDesc,
 			},
-			Code: out.Code,
+			Code: code,
 		}
 		module, err = ctx.ModuleInit(&moduleImpl)
 	}

diff --git a/py/util.go b/py/util.go
@@ -9,6 +9,10 @@ import (
 	"strconv"
 )
 
+var (
+	ErrUnsupportedObjType = errors.New("unsupported obj type")
+)
+
 func GetLen(obj Object) (Int, error) {
 	getlen, ok := obj.(I__len__)
 	if !ok {
@@ -53,10 +57,6 @@ func LoadTuple(args Tuple, vars []interface{}) error {
 	return nil
 }
 
-var (
-	ErrUnsupportedObjType = errors.New("unsupported obj type")
-)
-
 func LoadAttr(obj Object, attrName string, data interface{}) error {
 	attr, err := GetAttrString(obj, attrName)
 	if err != nil {