package runner import ( "context" "errors" "path/filepath" "runtime" "sync" "sync/atomic" "time" luaCtx "Moonshark/core/runner/context" "Moonshark/core/runner/sandbox" "Moonshark/core/utils/logger" luajit "git.sharkk.net/Sky/LuaJIT-to-Go" ) // Common errors var ( ErrRunnerClosed = errors.New("lua runner is closed") ErrInitFailed = errors.New("initialization failed") ErrStateNotReady = errors.New("lua state not ready") ErrTimeout = errors.New("operation timed out") ) // RunnerOption defines a functional option for configuring the Runner type RunnerOption func(*Runner) // State wraps a Lua state with its sandbox type State struct { L *luajit.State // The Lua state sandbox *sandbox.Sandbox // Associated sandbox index int // Index for debugging inUse bool // Whether the state is currently in use } // InitHook runs before executing a script type InitHook func(*luajit.State, *luaCtx.Context) error // FinalizeHook runs after executing a script type FinalizeHook func(*luajit.State, *luaCtx.Context, any) error // Runner runs Lua scripts using a pool of Lua states type Runner struct { states []*State // All states managed by this runner statePool chan int // Pool of available state indexes poolSize int // Size of the state pool moduleLoader *ModuleLoader // Module loader isRunning atomic.Bool // Whether the runner is active mu sync.RWMutex // Mutex for thread safety debug bool // Enable debug logging initHooks []InitHook // Hooks run before script execution finalizeHooks []FinalizeHook // Hooks run after script execution scriptDir string // Current script directory } // WithPoolSize sets the state pool size func WithPoolSize(size int) RunnerOption { return func(r *Runner) { if size > 0 { r.poolSize = size } } } // WithDebugEnabled enables debug output func WithDebugEnabled() RunnerOption { return func(r *Runner) { r.debug = true } } // WithLibDirs sets additional library directories func WithLibDirs(dirs ...string) RunnerOption { return func(r *Runner) { if r.moduleLoader == nil { r.moduleLoader = NewModuleLoader(&ModuleConfig{ LibDirs: dirs, }) } else { r.moduleLoader.config.LibDirs = dirs } } } // WithInitHook adds a hook to run before script execution func WithInitHook(hook InitHook) RunnerOption { return func(r *Runner) { r.initHooks = append(r.initHooks, hook) } } // WithFinalizeHook adds a hook to run after script execution func WithFinalizeHook(hook FinalizeHook) RunnerOption { return func(r *Runner) { r.finalizeHooks = append(r.finalizeHooks, hook) } } // NewRunner creates a new Runner with a pool of states func NewRunner(options ...RunnerOption) (*Runner, error) { // Default configuration runner := &Runner{ poolSize: runtime.GOMAXPROCS(0), debug: false, initHooks: make([]InitHook, 0, 4), finalizeHooks: make([]FinalizeHook, 0, 4), } // Apply options for _, opt := range options { opt(runner) } // Set up module loader if not already initialized if runner.moduleLoader == nil { config := &ModuleConfig{ ScriptDir: "", LibDirs: []string{}, } runner.moduleLoader = NewModuleLoader(config) } // Initialize states and pool runner.states = make([]*State, runner.poolSize) runner.statePool = make(chan int, runner.poolSize) // Create and initialize all states if err := runner.initializeStates(); err != nil { runner.Close() // Clean up already created states return nil, err } runner.isRunning.Store(true) return runner, nil } // debugLog logs a message if debug mode is enabled func (r *Runner) debugLog(format string, args ...interface{}) { if r.debug { logger.Debug("Runner "+format, args...) } } // initializeStates creates and initializes all states in the pool func (r *Runner) initializeStates() error { r.debugLog("is initializing %d states", r.poolSize) // Create all states for i := 0; i < r.poolSize; i++ { state, err := r.createState(i) if err != nil { return err } r.states[i] = state r.statePool <- i // Add index to the pool } return nil } // createState initializes a new Lua state func (r *Runner) createState(index int) (*State, error) { verbose := index == 0 if verbose { r.debugLog("Creating Lua state %d", index) } // Create a new state L := luajit.New() if L == nil { return nil, errors.New("failed to create Lua state") } // Create sandbox sb := sandbox.NewSandbox() if r.debug && verbose { sb.EnableDebug() } // Set up require system if err := r.moduleLoader.SetupRequire(L); err != nil { L.Cleanup() L.Close() return nil, ErrInitFailed } // Initialize all core modules from the registry if err := GlobalRegistry.Initialize(L, index); err != nil { L.Cleanup() L.Close() return nil, ErrInitFailed } // Set up sandbox after core modules are initialized if err := sb.Setup(L, index); err != nil { L.Cleanup() L.Close() return nil, ErrInitFailed } // Preload all modules if err := r.moduleLoader.PreloadModules(L); err != nil { L.Cleanup() L.Close() return nil, errors.New("failed to preload modules") } return &State{ L: L, sandbox: sb, index: index, inUse: false, }, nil } // Execute runs a script with context func (r *Runner) Execute(ctx context.Context, bytecode []byte, execCtx *luaCtx.Context, scriptPath string) (any, error) { if !r.isRunning.Load() { return nil, ErrRunnerClosed } // Set script directory if provided if scriptPath != "" { r.mu.Lock() r.scriptDir = filepath.Dir(scriptPath) r.moduleLoader.SetScriptDir(r.scriptDir) r.mu.Unlock() } // Get a state from the pool var stateIndex int select { case stateIndex = <-r.statePool: // Got a state case <-ctx.Done(): return nil, ctx.Err() case <-time.After(5 * time.Second): return nil, ErrTimeout } // Get the actual state state := r.states[stateIndex] if state == nil { r.statePool <- stateIndex return nil, ErrStateNotReady } // Mark state as in use state.inUse = true // Ensure state is returned to pool when done defer func() { state.inUse = false if r.isRunning.Load() { select { case r.statePool <- stateIndex: // State returned to pool default: // Pool is full or closed } } }() // Run init hooks for _, hook := range r.initHooks { if err := hook(state.L, execCtx); err != nil { return nil, err } } // Get context values var ctxValues map[string]any if execCtx != nil { ctxValues = execCtx.Values } // Execute in sandbox with optimized context handling var result any var err error if execCtx != nil && execCtx.RequestCtx != nil { // Use OptimizedExecute directly with the full context if we have RequestCtx result, err = state.sandbox.OptimizedExecute(state.L, bytecode, &luaCtx.Context{ Values: ctxValues, RequestCtx: execCtx.RequestCtx, }) } else { // Otherwise use standard Execute with just values result, err = state.sandbox.Execute(state.L, bytecode, ctxValues) } if err != nil { return nil, err } // Run finalize hooks for _, hook := range r.finalizeHooks { if hookErr := hook(state.L, execCtx, result); hookErr != nil { return nil, hookErr } } // Check for HTTP response if we don't have a RequestCtx or if we still have a result if execCtx == nil || execCtx.RequestCtx == nil || result != nil { httpResp, hasResponse := sandbox.GetHTTPResponse(state.L) if hasResponse { // Set result as body if not already set if httpResp.Body == nil { httpResp.Body = result } // Apply directly to request context if available if execCtx != nil && execCtx.RequestCtx != nil { sandbox.ApplyHTTPResponse(httpResp, execCtx.RequestCtx) sandbox.ReleaseResponse(httpResp) return nil, nil } return httpResp, nil } } return result, err } // Run executes a Lua script (convenience wrapper) func (r *Runner) Run(bytecode []byte, execCtx *luaCtx.Context, scriptPath string) (any, error) { return r.Execute(context.Background(), bytecode, execCtx, scriptPath) } // Close gracefully shuts down the Runner func (r *Runner) Close() error { r.mu.Lock() defer r.mu.Unlock() if !r.isRunning.Load() { return ErrRunnerClosed } r.isRunning.Store(false) // Drain the state pool for { select { case <-r.statePool: // Drain one state default: // Pool is empty goto cleanup } } cleanup: // Clean up all states for i, state := range r.states { if state != nil { state.L.Cleanup() state.L.Close() r.states[i] = nil } } return nil } // RefreshStates rebuilds all states in the pool func (r *Runner) RefreshStates() error { r.mu.Lock() defer r.mu.Unlock() if !r.isRunning.Load() { return ErrRunnerClosed } // Drain all states from the pool for { select { case <-r.statePool: // Drain one state default: // Pool is empty goto cleanup } } cleanup: // Destroy all existing states for i, state := range r.states { if state != nil { if state.inUse { r.debugLog("Warning: attempting to refresh state %d that is in use", i) } state.L.Cleanup() state.L.Close() r.states[i] = nil } } // Reinitialize all states if err := r.initializeStates(); err != nil { return err } r.debugLog("All states refreshed successfully") return nil } // AddInitHook adds a hook to be called before script execution func (r *Runner) AddInitHook(hook InitHook) { r.mu.Lock() defer r.mu.Unlock() r.initHooks = append(r.initHooks, hook) } // AddFinalizeHook adds a hook to be called after script execution func (r *Runner) AddFinalizeHook(hook FinalizeHook) { r.mu.Lock() defer r.mu.Unlock() r.finalizeHooks = append(r.finalizeHooks, hook) } // GetStateCount returns the number of initialized states func (r *Runner) GetStateCount() int { r.mu.RLock() defer r.mu.RUnlock() count := 0 for _, state := range r.states { if state != nil { count++ } } return count } // GetActiveStateCount returns the number of states currently in use func (r *Runner) GetActiveStateCount() int { r.mu.RLock() defer r.mu.RUnlock() count := 0 for _, state := range r.states { if state != nil && state.inUse { count++ } } return count } // GetModuleCount returns the number of loaded modules in the first available state func (r *Runner) GetModuleCount() int { r.mu.RLock() defer r.mu.RUnlock() if !r.isRunning.Load() { return 0 } // Find first available state for _, state := range r.states { if state != nil && !state.inUse { // Execute a Lua snippet to count modules if res, err := state.L.ExecuteWithResult(` local count = 0 for _ in pairs(package.loaded) do count = count + 1 end return count `); err == nil { if num, ok := res.(float64); ok { return int(num) } } break } } return 0 } // NotifyFileChanged alerts the runner about file changes func (r *Runner) NotifyFileChanged(filePath string) bool { r.debugLog("File change detected: %s", filePath) // Check if it's a module file module, isModule := r.moduleLoader.GetModuleByPath(filePath) if isModule { r.debugLog("File is a module: %s", module) return r.RefreshModule(module) } // For non-module files, refresh all states if err := r.RefreshStates(); err != nil { r.debugLog("Failed to refresh states: %v", err) return false } return true } // RefreshModule refreshes a specific module across all states func (r *Runner) RefreshModule(moduleName string) bool { r.mu.RLock() defer r.mu.RUnlock() if !r.isRunning.Load() { return false } r.debugLog("Refreshing module: %s", moduleName) // Check if it's a core module coreName, isCore := GlobalRegistry.MatchModuleName(moduleName) success := true for _, state := range r.states { if state == nil || state.inUse { continue } // Invalidate module in Lua if err := state.L.DoString(`package.loaded["` + moduleName + `"] = nil`); err != nil { success = false continue } // For core modules, reinitialize them if isCore { if err := GlobalRegistry.InitializeModule(state.L, coreName); err != nil { success = false } } } return success }