373 lines
9.0 KiB
Go
373 lines
9.0 KiB
Go
package workers
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
|
|
)
|
|
|
|
// Common errors
|
|
var (
|
|
ErrLoopClosed = errors.New("event loop is closed")
|
|
ErrExecutionTimeout = errors.New("script execution timed out")
|
|
)
|
|
|
|
// StateInitFunc is a function that initializes a Lua state
|
|
type StateInitFunc func(*luajit.State) error
|
|
|
|
// EventLoop represents a single-threaded Lua execution environment
|
|
type EventLoop struct {
|
|
state *luajit.State // Single Lua state for all executions
|
|
jobQueue chan job // Channel for receiving jobs
|
|
quit chan struct{} // Channel for shutdown signaling
|
|
wg sync.WaitGroup // WaitGroup for clean shutdown
|
|
isRunning atomic.Bool // Flag to track if loop is running
|
|
timeout time.Duration // Default timeout for script execution
|
|
stateInit StateInitFunc // Optional function to initialize Lua state
|
|
bufferSize int // Size of job queue buffer
|
|
}
|
|
|
|
// EventLoopConfig contains configuration options for creating an EventLoop
|
|
type EventLoopConfig struct {
|
|
// StateInit is a function to initialize the Lua state with custom modules and functions
|
|
StateInit StateInitFunc
|
|
|
|
// BufferSize is the size of the job queue buffer (default: 100)
|
|
BufferSize int
|
|
|
|
// Timeout is the default execution timeout (default: 30s, 0 means no timeout)
|
|
Timeout time.Duration
|
|
}
|
|
|
|
// NewEventLoop creates a new event loop with default configuration
|
|
func NewEventLoop() (*EventLoop, error) {
|
|
return NewEventLoopWithConfig(EventLoopConfig{})
|
|
}
|
|
|
|
// NewEventLoopWithInit creates a new event loop with a state initialization function
|
|
func NewEventLoopWithInit(init StateInitFunc) (*EventLoop, error) {
|
|
return NewEventLoopWithConfig(EventLoopConfig{
|
|
StateInit: init,
|
|
})
|
|
}
|
|
|
|
// NewEventLoopWithConfig creates a new event loop with custom configuration
|
|
func NewEventLoopWithConfig(config EventLoopConfig) (*EventLoop, error) {
|
|
// Set default values
|
|
bufferSize := config.BufferSize
|
|
if bufferSize <= 0 {
|
|
bufferSize = 100 // Default buffer size
|
|
}
|
|
|
|
timeout := config.Timeout
|
|
if timeout == 0 {
|
|
timeout = 30 * time.Second // Default timeout
|
|
}
|
|
|
|
// Initialize the Lua state
|
|
state := luajit.New()
|
|
if state == nil {
|
|
return nil, errors.New("failed to create Lua state")
|
|
}
|
|
|
|
// Create the event loop instance
|
|
el := &EventLoop{
|
|
state: state,
|
|
jobQueue: make(chan job, bufferSize),
|
|
quit: make(chan struct{}),
|
|
timeout: timeout,
|
|
stateInit: config.StateInit,
|
|
bufferSize: bufferSize,
|
|
}
|
|
el.isRunning.Store(true)
|
|
|
|
// Set up the sandbox environment
|
|
if err := setupSandbox(el.state); err != nil {
|
|
state.Close()
|
|
return nil, err
|
|
}
|
|
|
|
// Initialize the state if needed
|
|
if el.stateInit != nil {
|
|
if err := el.stateInit(el.state); err != nil {
|
|
state.Close()
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
// Start the event loop
|
|
el.wg.Add(1)
|
|
go el.run()
|
|
|
|
return el, nil
|
|
}
|
|
|
|
// run is the main event loop goroutine
|
|
func (el *EventLoop) run() {
|
|
defer el.wg.Done()
|
|
defer el.state.Close()
|
|
|
|
for {
|
|
select {
|
|
case job, ok := <-el.jobQueue:
|
|
if !ok {
|
|
// Job queue closed, exit
|
|
return
|
|
}
|
|
|
|
// Execute job with timeout if configured
|
|
if el.timeout > 0 {
|
|
el.executeJobWithTimeout(job)
|
|
} else {
|
|
// Execute without timeout
|
|
result := executeJobSandboxed(el.state, job)
|
|
job.Result <- result
|
|
}
|
|
|
|
case <-el.quit:
|
|
// Quit signal received, exit
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// executeJobWithTimeout executes a job with a timeout
|
|
func (el *EventLoop) executeJobWithTimeout(j job) {
|
|
// Create a context with timeout
|
|
ctx, cancel := context.WithTimeout(context.Background(), el.timeout)
|
|
defer cancel()
|
|
|
|
// Create a channel for the result
|
|
resultCh := make(chan JobResult, 1)
|
|
|
|
// Execute the job in a separate goroutine
|
|
go func() {
|
|
result := executeJobSandboxed(el.state, j)
|
|
select {
|
|
case resultCh <- result:
|
|
// Result sent successfully
|
|
case <-ctx.Done():
|
|
// Context canceled, result no longer needed
|
|
}
|
|
}()
|
|
|
|
// Wait for result or timeout
|
|
select {
|
|
case result := <-resultCh:
|
|
// Send result to the original channel
|
|
j.Result <- result
|
|
case <-ctx.Done():
|
|
// Timeout occurred
|
|
j.Result <- JobResult{nil, ErrExecutionTimeout}
|
|
// NOTE: The Lua execution continues in the background until it completes,
|
|
// but the result is discarded. This is a compromise to avoid forcibly
|
|
// terminating Lua code which could corrupt the state.
|
|
}
|
|
}
|
|
|
|
// Submit sends a job to the event loop
|
|
func (el *EventLoop) Submit(bytecode []byte, execCtx *Context) (any, error) {
|
|
return el.SubmitWithContext(context.Background(), bytecode, execCtx)
|
|
}
|
|
|
|
// SubmitWithTimeout sends a job to the event loop with a specific timeout
|
|
func (el *EventLoop) SubmitWithTimeout(bytecode []byte, execCtx *Context, timeout time.Duration) (any, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return el.SubmitWithContext(ctx, bytecode, execCtx)
|
|
}
|
|
|
|
// SubmitWithContext sends a job to the event loop with context for cancellation
|
|
func (el *EventLoop) SubmitWithContext(ctx context.Context, bytecode []byte, execCtx *Context) (any, error) {
|
|
if !el.isRunning.Load() {
|
|
return nil, ErrLoopClosed
|
|
}
|
|
|
|
resultChan := make(chan JobResult, 1)
|
|
j := job{
|
|
Bytecode: bytecode,
|
|
Context: execCtx,
|
|
Result: resultChan,
|
|
}
|
|
|
|
// Submit job with context
|
|
select {
|
|
case el.jobQueue <- j:
|
|
// Job submitted
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
|
|
// Wait for result with context
|
|
select {
|
|
case result := <-resultChan:
|
|
return result.Value, result.Error
|
|
case <-ctx.Done():
|
|
// Context canceled, but the job might still be processed
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
// SetTimeout updates the default timeout for script execution
|
|
func (el *EventLoop) SetTimeout(timeout time.Duration) {
|
|
el.timeout = timeout
|
|
}
|
|
|
|
// Shutdown gracefully shuts down the event loop
|
|
func (el *EventLoop) Shutdown() error {
|
|
if !el.isRunning.Load() {
|
|
return ErrLoopClosed
|
|
}
|
|
el.isRunning.Store(false)
|
|
|
|
// Signal event loop to quit
|
|
close(el.quit)
|
|
|
|
// Wait for event loop to finish
|
|
el.wg.Wait()
|
|
|
|
// Close job queue
|
|
close(el.jobQueue)
|
|
|
|
return nil
|
|
}
|
|
|
|
// setupSandbox initializes the sandbox environment in the Lua state
|
|
func setupSandbox(state *luajit.State) error {
|
|
// This is the Lua script that creates our sandbox function
|
|
setupScript := `
|
|
-- Create a function to run code in a sandbox environment
|
|
function __create_sandbox()
|
|
-- Create new environment table
|
|
local env = {}
|
|
|
|
-- Add standard library modules (can be restricted as needed)
|
|
env.string = string
|
|
env.table = table
|
|
env.math = math
|
|
env.os = {
|
|
time = os.time,
|
|
date = os.date,
|
|
difftime = os.difftime,
|
|
clock = os.clock
|
|
}
|
|
env.tonumber = tonumber
|
|
env.tostring = tostring
|
|
env.type = type
|
|
env.pairs = pairs
|
|
env.ipairs = ipairs
|
|
env.next = next
|
|
env.select = select
|
|
env.unpack = unpack
|
|
env.pcall = pcall
|
|
env.xpcall = xpcall
|
|
env.error = error
|
|
env.assert = assert
|
|
|
|
-- Allow access to package.loaded for modules
|
|
env.require = function(name)
|
|
return package.loaded[name]
|
|
end
|
|
|
|
-- Create metatable to restrict access to _G
|
|
local mt = {
|
|
__index = function(t, k)
|
|
-- First check in env table
|
|
local v = rawget(env, k)
|
|
if v ~= nil then return v end
|
|
|
|
-- If not found, check for registered modules/functions
|
|
local moduleValue = _G[k]
|
|
if type(moduleValue) == "table" or
|
|
type(moduleValue) == "function" then
|
|
return moduleValue
|
|
end
|
|
|
|
return nil
|
|
end,
|
|
__newindex = function(t, k, v)
|
|
rawset(env, k, v)
|
|
end
|
|
}
|
|
|
|
setmetatable(env, mt)
|
|
return env
|
|
end
|
|
|
|
-- Create function to execute code with a sandbox
|
|
function __run_sandboxed(f, ctx)
|
|
local env = __create_sandbox()
|
|
|
|
-- Add context to the environment if provided
|
|
if ctx then
|
|
env.ctx = ctx
|
|
end
|
|
|
|
-- Set the environment and run the function
|
|
setfenv(f, env)
|
|
return f()
|
|
end
|
|
`
|
|
|
|
return state.DoString(setupScript)
|
|
}
|
|
|
|
// executeJobSandboxed runs a script in a sandbox environment
|
|
func executeJobSandboxed(state *luajit.State, j job) JobResult {
|
|
// Set up context if provided
|
|
if j.Context != nil {
|
|
// Push context table
|
|
state.NewTable()
|
|
|
|
// Add values to context table
|
|
for key, value := range j.Context.Values {
|
|
// Push key
|
|
state.PushString(key)
|
|
|
|
// Push value
|
|
if err := state.PushValue(value); err != nil {
|
|
state.Pop(1) // Pop table
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Set table[key] = value
|
|
state.SetTable(-3)
|
|
}
|
|
} else {
|
|
// Push nil if no context
|
|
state.PushNil()
|
|
}
|
|
|
|
// Load bytecode
|
|
if err := state.LoadBytecode(j.Bytecode, "script"); err != nil {
|
|
state.Pop(1) // Pop context
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Get the sandbox runner function
|
|
state.GetGlobal("__run_sandboxed")
|
|
|
|
// Push loaded function and context as arguments
|
|
state.PushCopy(-2) // Copy the loaded function
|
|
state.PushCopy(-4) // Copy the context table or nil
|
|
|
|
// Remove the original function and context
|
|
state.Remove(-5) // Remove original context
|
|
state.Remove(-4) // Remove original function
|
|
|
|
// Call the sandbox runner with 2 args (function and context), expecting 1 result
|
|
if err := state.Call(2, 1); err != nil {
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Get result
|
|
value, err := state.ToValue(-1)
|
|
state.Pop(1) // Pop result
|
|
|
|
return JobResult{value, err}
|
|
}
|