improve stacktraces on errors

3 files changed, 84 insertions, 16 deletions

diff --git a/lua.go b/lua.go
index 43cd3f3..9be56a5 100644
--- a/lua.go
+++ b/lua.go
@@ -39,6 +39,7 @@ type Lua struct {
 	running **Lua
 	yield func()
 	resume func() bool
+	tracebackHandler LuaRef
 }
 
 //export luna_run_go_func
@@ -53,6 +54,10 @@ func (l *Lua) Start() {
 	l.l = C.luaL_newstate()
 	l.running = &l
 	C.luaL_openlibs(l.l)
+
+	// Put debug handler into the registry for later copies to the stack.
+	l.PushGoFunction(TracebackHandler)
+	l.tracebackHandler = l.PopToRef()
 }
 
 // Close closes the Lua context.
@@ -62,11 +67,12 @@ func (l *Lua) Close() {
 
 // Require loads and executes the file pushing results onto the Lua stack.
 func (l *Lua) Require(file string) error {
+	l.PushTracebackHandler()
 	err := l.LoadFile(file)
 	if err != nil {
 		return errors.New("could not open the file:\n" + err.Error())
 	}
-	err = l.PCall(0, C.LUA_MULTRET)
+	err = l.PCall(0, C.LUA_MULTRET, -2)
 	if err != nil {
 		return errors.New("could not execute the file:\n" + err.Error())
 	}
@@ -75,9 +81,11 @@ func (l *Lua) Require(file string) error {
 
 // PCall calls a function with the stack index (-nargs-1) expecting nresults
 // number of results.
-func (l *Lua) PCall(nargs int, nresults int) error {
+func (l *Lua) PCall(nargs int, nresults int, errfunc int) error {
 	*l.running = l
-	res := C.lua_pcallk(l.l, C.int(nargs), C.int(nresults), 0, 0, nil)
+	res := C.lua_pcallk(
+		l.l, C.int(nargs), C.int(nresults), C.int(errfunc), 0, nil,
+	)
 	if res != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
@@ -187,6 +195,14 @@ func (l *Lua) SetTableItem(key string) {
 	C.lua_setfield(l.l, -2, C.CString(key))
 }
 
+// GetTableItem gets a value from the table (offset -1) under the specified key
+// and puts it onto the stack.
+func (l *Lua) GetTableItem(key string) {
+	cstr := C.CString(key)
+	defer C.free(unsafe.Pointer(cstr))
+	C.lua_getfield(l.l, -1, C.CString(key))
+}
+
 // PushAny pushes value v onto the stack.
 func (l *Lua) PushAny(v any) error {
 	switch v.(type) {
@@ -263,6 +279,11 @@ func (l *Lua) PushFromRef(ref LuaRef) {
 	C.lua_rawgeti(l.l, C.LUA_REGISTRYINDEX, C.longlong(ref));
 }
 
+// Unref removes the ref from the registry.
+func (l *Lua) Unref(ref LuaRef) {
+	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+}
+
 // Type returns type of the value sitting at n index on the stack.
 func (l *Lua) Type(n int) int {
 	return int(C.lua_type(l.l, C.int(n)))
@@ -309,12 +330,13 @@ func (l *Lua) Next() bool {
 func (l *Lua) LoadString(code string) error {
 	cstr := C.CString(code)
 	defer C.free(unsafe.Pointer(cstr))
+	l.PushTracebackHandler()
 	if C.luaL_loadstring(l.l, cstr) != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
 		return errors.New(errMsg)
 	}
-	err := l.PCall(0, 0)
+	err := l.PCall(0, 0, -2)
 	if err != nil {
 		return err
 	}
@@ -469,9 +491,22 @@ func (l *Lua) NewThread(yield func(), resume func() bool) *Lua {
 		running: l.running,
 		resume: resume,
 		yield: yield,
+		tracebackHandler: l.tracebackHandler,
 	}
 }
 
-func (l *Lua) Unref(ref LuaRef) {
-	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+// PushTracebackHandler puts the traceback handler onto the stack for using
+// as a errfunc argument when calling PCall.
+func (l *Lua) PushTracebackHandler() {
+	l.PushFromRef(l.tracebackHandler)
+}
+
+// TracebackHandler handles stack trace formatting on errors.
+func TracebackHandler(l *Lua) int {
+	msg := l.ToString(-1)
+	l.GetGlobal("debug")
+	l.GetTableItem("traceback")
+	l.PushString(msg)
+	C.lua_callk(l.l, 1, 1, 0, nil)
+	return 1
 }
diff --git a/luajit.go b/luajit.go
index 9414787..7439023 100644
--- a/luajit.go
+++ b/luajit.go
@@ -39,6 +39,7 @@ type Lua struct {
 	running **Lua
 	yield func()
 	resume func() bool
+	tracebackHandler LuaRef
 }
 
 //export luna_run_go_func
@@ -53,6 +54,10 @@ func (l *Lua) Start() {
 	l.l = C.luaL_newstate()
 	l.running = &l
 	C.luaL_openlibs(l.l)
+
+	// Put debug handler into the registry for later copies to the stack.
+	l.PushGoFunction(TracebackHandler)
+	l.tracebackHandler = l.PopToRef()
 }
 
 // Close closes the Lua context.
@@ -62,12 +67,13 @@ func (l *Lua) Close() {
 
 // Require loads and executes the file pushing results onto the Lua stack.
 func (l *Lua) Require(file string) error {
+	l.PushTracebackHandler()
 	if !l.LoadFile(file) {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
 		return errors.New("could not open the file:\n" + errMsg)
 	}
-	err := l.PCall(0, C.LUA_MULTRET)
+	err := l.PCall(0, C.LUA_MULTRET, -2)
 	if err != nil {
 		return errors.New("could not execute the file:\n" + err.Error())
 	}
@@ -76,9 +82,9 @@ func (l *Lua) Require(file string) error {
 
 // PCall calls a function with the stack index (-nargs-1) expecting nresults
 // number of results.
-func (l *Lua) PCall(nargs int, nresults int) error {
+func (l *Lua) PCall(nargs int, nresults int, errfunc int) error {
 	*l.running = l
-	res := C.lua_pcall(l.l, C.int(nargs), C.int(nresults), 0)
+	res := C.lua_pcall(l.l, C.int(nargs), C.int(nresults), C.int(errfunc))
 	if res != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
@@ -183,6 +189,14 @@ func (l *Lua) SetTableItem(key string) {
 	C.lua_setfield(l.l, -2, C.CString(key))
 }
 
+// GetTableItem gets a value from the table (offset -1) under the specified key
+// and puts it onto the stack.
+func (l *Lua) GetTableItem(key string) {
+	cstr := C.CString(key)
+	defer C.free(unsafe.Pointer(cstr))
+	C.lua_getfield(l.l, -1, C.CString(key))
+}
+
 // PushAny pushes value v onto the stack.
 func (l *Lua) PushAny(v any) error {
 	switch v.(type) {
@@ -256,6 +270,11 @@ func (l *Lua) PushFromRef(ref LuaRef) {
 	C.lua_rawgeti(l.l, C.LUA_REGISTRYINDEX, C.int(ref));
 }
 
+// Unref removes the ref from the registry.
+func (l *Lua) Unref(ref LuaRef) {
+	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+}
+
 // Type returns type of the value sitting at n index on the stack.
 func (l *Lua) Type(n int) int {
 	return int(C.lua_type(l.l, C.int(n)))
@@ -301,12 +320,13 @@ func (l *Lua) Next() bool {
 func (l *Lua) LoadString(code string) error {
 	cstr := C.CString(code)
 	defer C.free(unsafe.Pointer(cstr))
+	l.PushTracebackHandler()
 	if C.luaL_loadstring(l.l, cstr) != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
 		return errors.New(errMsg)
 	}
-	err := l.PCall(0, 0)
+	err := l.PCall(0, 0, -2)
 	if err != nil {
 		return err
 	}
@@ -461,9 +481,22 @@ func (l *Lua) NewThread(yield func(), resume func() bool) *Lua {
 		running: l.running,
 		resume: resume,
 		yield: yield,
+		tracebackHandler: l.tracebackHandler,
 	}
 }
 
-func (l *Lua) Unref(ref LuaRef) {
-	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+// PushTracebackHandler puts the traceback handler onto the stack for using
+// as a errfunc argument when calling PCall.
+func (l *Lua) PushTracebackHandler() {
+	l.PushFromRef(l.tracebackHandler)
+}
+
+// TracebackHandler handles stack trace formatting on errors.
+func TracebackHandler(l *Lua) int {
+	msg := l.ToString(-1)
+	l.GetGlobal("debug")
+	l.GetTableItem("traceback")
+	l.PushString(msg)
+	C.lua_call(l.l, 1, 1)
+	return 1
 }
diff --git a/worker.go b/worker.go
index aa18b62..5ae69c2 100644
--- a/worker.go
+++ b/worker.go
@@ -89,10 +89,10 @@ func (w *Worker) Start(argv []string, module map[string]any) error {
 	w.lua.SetGlobal("luna")
 
 	err = w.lua.Require(argv[0])
-
 	if err != nil {
 		return err
 	}
+
 	w.started = true
 	return nil
 }
@@ -123,6 +123,7 @@ func (w *Worker) Listen(queue chan *HTTPRequest) {
 			return
 		}
 
+		l.PushTracebackHandler()
 		l.PushFromRef(w.routes[r.route])
 		res := make(map[string]any)
 		res["method"] = r.request.Method
@@ -164,15 +165,14 @@ func (w *Worker) Listen(queue chan *HTTPRequest) {
 			return
 		}
 
-		err = l.PCall(1, 3)
+		err = l.PCall(1, 3, -3)
 		if err != nil {
 			r.result <- &HTTPResponse{
 				Code: 500,
 				Headers: make(map[string]string),
 				Body: "server error",
 			}
-			log.Println("could not process request:", err)
-			// TODO: print lua stack as well?
+			log.Println("could not process request:\n" + err.Error())
 			return
 		}