Moonshark/core/workers/worker.go

package workers

import (
	"context"
	"errors"
	"sync/atomic"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// Common errors
var (
	ErrPoolClosed = errors.New("worker pool is closed")
	ErrNoWorkers  = errors.New("no workers available")
)

// worker represents a single Lua execution worker
type worker struct {
	pool  *Pool         // Reference to the pool
	state *luajit.State // Lua state
	id    uint32        // Worker ID
}

// run is the main worker function that processes jobs
func (w *worker) run() {
	defer w.pool.wg.Done()

	// Initialize Lua state
	w.state = luajit.New()
	if w.state == nil {
		// Worker failed to initialize, decrement counter
		atomic.AddUint32(&w.pool.workers, ^uint32(0))
		return
	}
	defer w.state.Close()

	// Set up reset function for clearing state between requests
	if err := w.setupResetFunction(); err != nil {
		// Worker failed to initialize reset function, decrement counter
		atomic.AddUint32(&w.pool.workers, ^uint32(0))
		return
	}

	// Main worker loop
	for {
		select {
		case job, ok := <-w.pool.jobs:
			if !ok {
				// Jobs channel closed, exit
				return
			}

			// Execute job
			result := w.executeJob(job)
			job.Result <- result

		case <-w.pool.quit:
			// Quit signal received, exit
			return
		}
	}
}

// setupResetFunction initializes the reset function for clearing globals
func (w *worker) setupResetFunction() error {
	resetScript := `
		-- Create reset function to efficiently clear globals after each request
		function __reset_globals()
			-- Only keep builtin globals, remove all user-defined globals
			local preserve = {
				["_G"] = true, ["_VERSION"] = true, ["__reset_globals"] = true,
				["assert"] = true, ["collectgarbage"] = true, ["coroutine"] = true,
				["debug"] = true, ["dofile"] = true, ["error"] = true,
				["getmetatable"] = true, ["io"] = true, ["ipairs"] = true,
				["load"] = true, ["loadfile"] = true, ["loadstring"] = true, 
				["math"] = true, ["next"] = true, ["os"] = true,
				["package"] = true, ["pairs"] = true, ["pcall"] = true,
				["print"] = true, ["rawequal"] = true, ["rawget"] = true,
				["rawset"] = true, ["require"] = true, ["select"] = true,
				["setmetatable"] = true, ["string"] = true, ["table"] = true,
				["tonumber"] = true, ["tostring"] = true, ["type"] = true,
				["unpack"] = true, ["xpcall"] = true
			}
			
			-- Clear all non-standard globals
			for name in pairs(_G) do
				if not preserve[name] then
					_G[name] = nil
				end
			end
			
			-- Run garbage collection to release memory
			collectgarbage('collect')
		end
	`

	return w.state.DoString(resetScript)
}

// resetState prepares the Lua state for a new job
func (w *worker) resetState() {
	w.state.DoString("__reset_globals()")
}

// setParams sets job parameters as a global 'params' table
func (w *worker) setParams(params map[string]interface{}) error {
	// Create new table for params
	w.state.NewTable()

	// Add each parameter to the table
	for key, value := range params {
		// Push key
		w.state.PushString(key)

		// Push value
		if err := w.state.PushValue(value); err != nil {
			return err
		}

		// Set table[key] = value
		w.state.SetTable(-3)
	}

	// Set the table as global 'params'
	w.state.SetGlobal("params")

	return nil
}

// executeJob executes a Lua job in the worker's state
func (w *worker) executeJob(j job) JobResult {
	// Reset state before execution
	w.resetState()

	// Set parameters
	if j.Params != nil {
		if err := w.setParams(j.Params); err != nil {
			return JobResult{nil, err}
		}
	}

	// Load bytecode
	if err := w.state.LoadBytecode(j.Bytecode, "script"); err != nil {
		return JobResult{nil, err}
	}

	// Execute script with one result
	if err := w.state.RunBytecodeWithResults(1); err != nil {
		return JobResult{nil, err}
	}

	// Get result
	value, err := w.state.ToValue(-1)
	w.state.Pop(1) // Pop result

	return JobResult{value, err}
}

// Submit sends a job to the worker pool
func (p *Pool) Submit(bytecode []byte, params map[string]interface{}) (interface{}, error) {
	return p.SubmitWithContext(context.Background(), bytecode, params)
}