lstate_pool.go

package pool

// StateMutator will modify an LState before it goes into the pool. This can
// run any number of scripts as necessary such as registring libraries,
// executing code, etc...
type StateMutator func(*lua.LState)

// PooledState wraps a Lua state. It's purpose is provide a means with which
// to return the state to the StatePool when it's not longer being used.
type PooledState struct {
	*lua.LState
	pool *StatePool
}

// Release will push the state back into the queue for available states for
// the current PooledState as well as nil out the reference to the state
// to prevent continued usage of the state.
func (pe *PooledState) Release() {
	if pe.LState != nil {
		pe.pool.states <- pe.LState
		pe.LState = nil
	}
}

// StatePool represents a grouping of predefined/preloaded state that can be
// grabbed for use when Lua scripts need to run.
type StatePool struct {
	MaxPoolSize uint8
	mutatorFn   StateMutator
	numStates   uint8
	states      chan *lua.LState
	mutex       *sync.Mutex
}

// NewStatePool constructs a new pool with the specific maximum size and the
// state mutator. It will seed the pool with one state.
func NewStatePool(poolSize uint8, mutator StateMutator) *StatePool {
	if poolSize == 0 {
		poolSize = 1
	}
	ep := &StatePool{
		MaxPoolSize: poolSize,
		mutatorFn:   mutator,
		numStates:   1,
		states:      make(chan *lua.LState, poolSize),
		mutex:       new(sync.Mutex),
	}
	ep.states <- ep.generateState()

	return ep
}

// Get will fetch the next available state from the StatePool. If no states
// are available and the maximum number of active states in the pool have been
// created yet then the spawner will be invoked to spawn a new state and return
// that.
func (ep *StatePool) Get() *PooledState {
	var state *lua.LState
	if len(ep.states) > 0 {
		state = <-ep.states
	} else if ep.numStates < ep.MaxPoolSize {
		ep.mutex.Lock()
		state = ep.generateState()
		ep.numStates++
		ep.mutex.Unlock()
	} else {
		state = <-ep.states
	}

	pe := &PooledState{
		LState: state,
		pool:   ep,
	}
	// NOTE: precaution to prevent leaks for long running servers, not a perfect
	//       solution. BE DILIGENT AND RELEASE YOUR POOLED STATES!!
	runtime.SetFinalizer(pe, (*PooledState).Release)

	return pe
}

func (ep *StatePool) generateState() *lua.LState {
	state := lua.NewState()
	ep.mutatorFn(state)

	return state
}