lua: Init with config and pool (#321)

* lua: Intial impl with config and pool

Signed-off-by: James Hillyerd <james@hillyerd.com>

pkg/extension/luahost/lua.go

@@ -0,0 +1,85 @@
+package luahost
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/inbucket/inbucket/pkg/config"
+	"github.com/inbucket/inbucket/pkg/extension"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	lua "github.com/yuin/gopher-lua"
+	"github.com/yuin/gopher-lua/parse"
+)
+
+// Host of Lua extensions.
+type Host struct {
+	Functions  []string // Functions detected in lua script.
+	extHost    *extension.Host
+	pool       *statePool
+	logContext zerolog.Context
+}
+
+// New constructs a new Lua Host, pre-compiling the source.
+func New(conf config.Lua, extHost *extension.Host) (*Host, error) {
+	scriptPath := conf.Path
+	if scriptPath == "" {
+		return nil, nil
+	}
+
+	logContext := log.With().Str("module", "lua")
+	logger := logContext.Str("phase", "startup").Str("path", scriptPath).Logger()
+
+	// Pre-load, parse, and compile script.
+	if fi, err := os.Stat(scriptPath); err != nil {
+		logger.Info().Msg("Script file not found")
+		return nil, nil
+	} else if fi.IsDir() {
+		return nil, fmt.Errorf("Lua script %v is a directory", scriptPath)
+	}
+
+	logger.Info().Msg("Loading script")
+	file, err := os.Open(scriptPath)
+	defer file.Close()
+	if err != nil {
+		return nil, err
+	}
+
+	return NewFromReader(extHost, bufio.NewReader(file), scriptPath)
+}
+
+// NewFromReader constructs a new Lua Host, loading Lua source from the provided reader.
+// The provided path is used in logging and error messages.
+func NewFromReader(extHost *extension.Host, r io.Reader, path string) (*Host, error) {
+	logContext := log.With().Str("module", "lua")
+
+	// Pre-parse, and compile script.
+	chunk, err := parse.Parse(r, path)
+	if err != nil {
+		return nil, err
+	}
+	proto, err := lua.Compile(chunk, path)
+	if err != nil {
+		return nil, err
+	}
+
+	// Build the pool and confirm LState is retrievable.
+	pool := newStatePool(proto)
+	h := &Host{extHost: extHost, pool: pool, logContext: logContext}
+	if ls, err := pool.getState(); err == nil {
+		// State creation works, put it back.
+		pool.putState(ls)
+	} else {
+		return nil, err
+	}
+
+	return h, nil
+}
+
+// CreateChannel creates a channel and places it into the named global variable
+// in newly created LStates.
+func (h *Host) CreateChannel(name string) chan lua.LValue {
+	return h.pool.createChannel(name)
+}

pkg/extension/luahost/lua_test.go

@@ -0,0 +1,18 @@
+package luahost_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/inbucket/inbucket/pkg/extension"
+	"github.com/inbucket/inbucket/pkg/extension/luahost"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEmptyScript(t *testing.T) {
+	script := ""
+	extHost := extension.NewHost()
+
+	_, err := luahost.NewFromReader(extHost, strings.NewReader(script), "test.lua")
+	require.NoError(t, err)
+}

pkg/extension/luahost/pool.go

@@ -0,0 +1,91 @@
+package luahost
+
+import (
+	"sync"
+
+	lua "github.com/yuin/gopher-lua"
+)
+
+type statePool struct {
+	sync.Mutex
+	funcProto *lua.FunctionProto         // Compiled lua.
+	states    []*lua.LState              // Pool of available LStates.
+	channels  map[string]chan lua.LValue // Global interop channels.
+}
+
+func newStatePool(funcProto *lua.FunctionProto) *statePool {
+	return &statePool{
+		funcProto: funcProto,
+		channels:  make(map[string]chan lua.LValue),
+	}
+}
+
+// newState creates a new LState and configures it. Lock must be held.
+func (lp *statePool) newState() (*lua.LState, error) {
+	ls := lua.NewState()
+
+	// Setup channels.
+	for name, ch := range lp.channels {
+		ls.SetGlobal(name, lua.LChannel(ch))
+	}
+
+	// Run compiled script.
+	ls.Push(ls.NewFunctionFromProto(lp.funcProto))
+	if err := ls.PCall(0, lua.MultRet, nil); err != nil {
+		return nil, err
+	}
+
+	return ls, nil
+}
+
+// getState returns a free LState, or creates a new one.
+func (lp *statePool) getState() (*lua.LState, error) {
+	lp.Lock()
+	defer lp.Unlock()
+
+	ln := len(lp.states)
+	if ln == 0 {
+		return lp.newState()
+	}
+
+	state := lp.states[ln-1]
+	lp.states = lp.states[0 : ln-1]
+
+	return state, nil
+}
+
+// putState returns the LState to the pool.
+func (lp *statePool) putState(state *lua.LState) {
+	if state.IsClosed() {
+		return
+	}
+
+	// Clear stack.
+	state.Pop(state.GetTop())
+
+	lp.Lock()
+	defer lp.Unlock()
+
+	lp.states = append(lp.states, state)
+}
+
+// createChannel creates a new channel, which will become a global variable in
+// newly created LStates.  We also destroy any pooled states.
+//
+// Warning: There may still be checked out LStates that will not have the value
+// set, which could be put back into the pool.
+func (lp *statePool) createChannel(name string) chan lua.LValue {
+	lp.Lock()
+	defer lp.Unlock()
+
+	ch := make(chan lua.LValue, 10)
+	lp.channels[name] = ch
+
+	// Flush state pool.
+	for _, s := range lp.states {
+		s.Close()
+	}
+	lp.states = lp.states[:0]
+
+	return ch
+}

pkg/extension/luahost/pool_test.go

@@ -0,0 +1,101 @@
+package luahost
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	lua "github.com/yuin/gopher-lua"
+	"github.com/yuin/gopher-lua/parse"
+)
+
+func makeEmptyPool() *statePool {
+	source := strings.NewReader("-- Empty source")
+
+	chunk, err := parse.Parse(source, "from string")
+	if err != nil {
+		panic(err)
+	}
+
+	proto, err := lua.Compile(chunk, "from string")
+	if err != nil {
+		panic(err)
+	}
+
+	return newStatePool(proto)
+}
+
+func TestPoolGetsDistinct(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	b, err := pool.getState()
+	require.NoError(t, err)
+
+	if a == b {
+		t.Error("Got pool a == b, expected distinct pools")
+	}
+}
+
+func TestPoolGrowsWithPuts(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	b, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	pool.putState(a)
+	pool.putState(b)
+
+	want := 2
+	if got := len(pool.states); got != want {
+		t.Errorf("len pool.states got %v, want %v", got, want)
+	}
+}
+
+// Closed LStates should not be added to the pool.
+func TestPoolPutDiscardsClosed(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	a.Close()
+	pool.putState(a)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to remain empty")
+}
+
+func TestPoolPutClearsStack(t *testing.T) {
+	pool := makeEmptyPool()
+
+	ls, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	// Setup stack.
+	ls.Push(lua.LNumber(4))
+	ls.Push(lua.LString("bacon"))
+	require.Equal(t, 2, ls.GetTop(), "Want stack to have two items")
+
+	// Return and verify stack cleared.
+	pool.putState(ls)
+	assert.Equal(t, 1, len(pool.states), "Wanted pool to have one item")
+	require.Equal(t, 0, ls.GetTop(), "Want stack to be empty")
+}
+
+func TestPoolSetsChannels(t *testing.T) {
+	pool := makeEmptyPool()
+	pool.createChannel("test_chan")
+
+	s, err := pool.getState()
+	require.NoError(t, err)
+
+	got := s.GetGlobal("test_chan")
+	assert.Equal(t, lua.LTChannel, got.Type(),
+		"Got global type %v, wanted LTChannel", got.Type().String())
+}