Support stacktraces

Author: cschleiden

internal/sync/coroutine.go

@@ -11,6 +11,10 @@ import (
 
 const DeadlockDetection = 40 * time.Second
 
+type CoroutineCreator interface {
+	NewCoroutine(ctx Context, fn func(Context) error)
+}
+
 type Coroutine interface {
 	// Execute continues execution of a blocked corouting and waits until
 	// it is finished or blocked again
@@ -28,7 +32,9 @@ type Coroutine interface {
 
 	Error() error
 
-	SetScheduler(s Scheduler)
+	SetCoroutineCreator(creator CoroutineCreator)
+
+	SetPanicHandler(handler func(interface{}) error)
 }
 
 type key int
@@ -47,13 +53,14 @@ type coState struct {
 	shouldExit atomic.Value // coroutine should exit
 	progress   atomic.Value // did the coroutine make progress since last yield?
 
-	err error
+	err          error
+	panicHandler func(interface{}) error
 
 	logger logger
 
 	deadlockDetection time.Duration
 
-	scheduler Scheduler
+	creator CoroutineCreator
 }
 
 func NewCoroutine(ctx Context, fn func(ctx Context) error) Coroutine {
@@ -64,7 +71,11 @@ func NewCoroutine(ctx Context, fn func(ctx Context) error) Coroutine {
 		defer s.finish() // Ensure we always mark the coroutine as finished
 		defer func() {
 			if r := recover(); r != nil {
-				s.err = fmt.Errorf("panic: %v", r)
+				if s.panicHandler != nil {
+					s.err = s.panicHandler(r)
+				} else {
+					s.err = fmt.Errorf("panic: %v", r)
+				}
 			}
 		}()
 
@@ -99,8 +110,12 @@ func (s *coState) finish() {
 	s.logger.Println("finish")
 }
 
-func (s *coState) SetScheduler(scheduler Scheduler) {
-	s.scheduler = scheduler
+func (s *coState) SetCoroutineCreator(creator CoroutineCreator) {
+	s.creator = creator
+}
+
+func (s *coState) SetPanicHandler(handler func(interface{}) error) {
+	s.panicHandler = handler
 }
 
 func (s *coState) Finished() bool {

internal/sync/coroutine_test.go

@@ -2,6 +2,7 @@ package sync
 
 import (
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
@@ -177,3 +178,18 @@ func Test_Coroutine_Panic(t *testing.T) {
 	require.Error(t, c.Error())
 	require.Equal(t, c.Error().Error(), "panic: test panic")
 }
+
+func Test_Coroutine_CallPanicHandlerIfSet(t *testing.T) {
+	c := NewCoroutine(Background(), func(ctx Context) error {
+		panic("test panic")
+	})
+
+	c.SetPanicHandler(func(r interface{}) error {
+		return fmt.Errorf("handled panic: %v", r)
+	})
+
+	c.Execute()
+
+	require.True(t, c.Finished())
+	require.EqualError(t, c.Error(), "handled panic: test panic")
+}

internal/sync/go.go

@@ -3,7 +3,7 @@ package sync
 func Go(ctx Context, f func(ctx Context)) {
 	cs := getCoState(ctx)
 
-	cs.scheduler.NewCoroutine(ctx, func(ctx Context) error {
+	cs.creator.NewCoroutine(ctx, func(ctx Context) error {
 		f(ctx)
 
 		return nil

internal/sync/scheduler.go

@@ -1,34 +1,32 @@
 package sync
 
-type Scheduler interface {
-	// Starts a new co-routine and tracks it in this scheduler
-	NewCoroutine(ctx Context, fn func(Context) error)
-
-	// Execute executes all coroutines until they are all blocked
-	Execute() error
-
-	RunningCoroutines() int
-
-	Exit()
-}
-
-type scheduler struct {
-	coroutines []Coroutine
+type Scheduler struct {
+	coroutines   []Coroutine
+	panicHandler func(interface{}) error
 }
 
-func NewScheduler() Scheduler {
-	return &scheduler{
+func NewScheduler() *Scheduler {
+	return &Scheduler{
 		coroutines: make([]Coroutine, 0),
 	}
 }
 
-func (s *scheduler) NewCoroutine(ctx Context, fn func(Context) error) {
+func (s *Scheduler) SetPanicHandler(handler func(interface{}) error) {
+	s.panicHandler = handler
+}
+
+// Starts a new co-routine and tracks it in this scheduler
+func (s *Scheduler) NewCoroutine(ctx Context, fn func(Context) error) {
 	c := NewCoroutine(ctx, fn)
+	if s.panicHandler != nil {
+		c.SetPanicHandler(s.panicHandler)
+	}
 	s.coroutines = append(s.coroutines, c)
-	c.SetScheduler(s)
+	c.SetCoroutineCreator(s)
 }
 
-func (s *scheduler) Execute() error {
+// Execute executes all coroutines until they are all blocked
+func (s *Scheduler) Execute() error {
 	allBlocked := false
 	for !allBlocked {
 		allBlocked = true
@@ -60,11 +58,11 @@ func (s *scheduler) Execute() error {
 	return nil
 }
 
-func (s *scheduler) RunningCoroutines() int {
+func (s *Scheduler) RunningCoroutines() int {
 	return len(s.coroutines)
 }
 
-func (s *scheduler) Exit() {
+func (s *Scheduler) Exit() {
 	for _, c := range s.coroutines {
 		c.Exit()
 	}

internal/workflow/workflow.go

@@ -14,14 +14,17 @@ import (
 type Workflow interface{}
 
 type workflow struct {
-	s      sync.Scheduler
+	s      *sync.Scheduler
 	fn     reflect.Value
 	result payload.Payload
 	err    error
 }
 
 func NewWorkflow(workflowFn reflect.Value) *workflow {
 	s := sync.NewScheduler()
+	s.SetPanicHandler(func(i interface{}) error {
+		return workflowerrors.
+	})
 
 	return &workflow{
 		s:  s,

internal/workflowerrors/error.go

@@ -1,14 +1,34 @@
 package workflowerrors
 
-import "errors"
+import (
+	"encoding/json"
+	"errors"
+)
 
 type Error struct {
 	Type    string `json:"type,omitempty"`
 	Message string `json:"message,omitempty"`
 
 	Permanent  bool   `json:"permanent,omitempty"`
 	Cause      error  `json:"cause,omitempty"`
-	StackTrace string `json:"stackTrace,omitempty"`
+	Stacktrace string `json:"stacktrace,omitempty"`
+}
+
+func (e *Error) UnmarshalJSON(b []byte) error {
+	type Alias Error
+	a := &struct {
+		Cause *Error `json:"cause,omitempty"`
+		*Alias
+	}{}
+
+	if err := json.Unmarshal(b, &a); err != nil {
+		return err
+	}
+
+	*e = *(*Error)(a.Alias)
+	e.Cause = a.Cause
+
+	return nil
 }
 
 func (we *Error) Error() string {
@@ -19,6 +39,10 @@ func (we *Error) Unwrap() error {
 	return we.Cause
 }
 
+func (we *Error) Stack() string {
+	return we.Stacktrace
+}
+
 var _ error = (*Error)(nil)
 
 // FromError wraps the given error into a workflow error which can be persisted and restored
@@ -37,6 +61,10 @@ func FromError(err error) *Error {
 		Message: err.Error(),
 	}
 
+	if stackTracer, ok := err.(interface{ Stack() string }); ok {
+		e.Stacktrace = stackTracer.Stack()
+	}
+
 	if cause := errors.Unwrap(err); cause != nil {
 		e.Cause = FromError(cause)
 	}
@@ -55,7 +83,7 @@ func ToError(err *Error) error {
 
 	switch err.Type {
 	case getErrorType(&PanicError{}):
-		return &PanicError{message: e.Message, stacktrace: e.StackTrace}
+		return &PanicError{message: e.Message, stacktrace: e.Stacktrace}
 
 	default:
 		// Keep *Error

internal/workflowerrors/panic.go

@@ -9,12 +9,13 @@ func (pe *PanicError) Error() string {
 	return pe.message
 }
 
-func (pe *PanicError) Stacktrace() string {
+func (pe *PanicError) Stack() string {
 	return pe.stacktrace
 }
 
-func NewPanicError(msg string) error {
+func NewPanicError(msg string) *PanicError {
 	return &PanicError{
-		message: msg,
+		message:    msg,
+		stacktrace: stack(3), // Skip new panic error and immediate caller
 	}
 }

internal/workflowerrors/panic_test.go

@@ -0,0 +1,16 @@
+package workflowerrors
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_NewPanicError(t *testing.T) {
+	e := func() *PanicError {
+		return NewPanicError("test")
+	}()
+
+	require.NotContains(t, e.Stack(), "Test_NewPanicError.func1")
+	require.NotContains(t, e.Stack(), "NewPanicError")
+}

internal/workflowerrors/stack.go

@@ -1,8 +1,34 @@
 package workflowerrors
 
-import goerrors "github.com/go-errors/errors"
+import (
+	"bytes"
+	"runtime"
 
-func stack(err error) string {
-	goerr := goerrors.New(err)
-	return string(goerr.Stack())
+	goerrors "github.com/go-errors/errors"
+)
+
+const MaxStackDepth = 50
+
+// stack returns a stacktrace formatted via go-errors/errors
+func stack(skip int) string {
+	// get stack
+	stack := make([]uintptr, MaxStackDepth)
+	length := runtime.Callers(2+skip, stack[:])
+
+	// trim
+	stack = stack[:length]
+
+	frames := make([]goerrors.StackFrame, len(stack))
+	for i, pc := range stack {
+		frames[i] = goerrors.NewStackFrame(pc)
+	}
+
+	// Convert frames to trace
+	buf := bytes.Buffer{}
+
+	for _, frame := range frames {
+		buf.WriteString(frame.String())
+	}
+
+	return string(buf.Bytes())
 }

internal/workflowerrors/stack_test.go

@@ -0,0 +1,24 @@
+package workflowerrors
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_stack(t *testing.T) {
+	fn := func() {
+		s := stack(1)
+		require.NotContains(t, s, "Test_stack.func1")
+	}
+
+	foo(fn)
+}
+
+func foo(fn func()) {
+	bar(fn)
+}
+
+func bar(fn func()) {
+	fn()
+}