Goroutines global panic handler (#188)

* New global panic handler allowing to fail running tests and flush the recorder buffer before process exit

* Panic handler test

* Test race condition fix

* Panic handler test in scope

Author: tonyredondo

agent/agent.go

@@ -396,6 +396,15 @@ func (a *Agent) Stop() {
 	a.PrintReport()
 }
 
+// Flush agent buffer
+func (a *Agent) Flush() {
+	a.logger.Println("Flushing agent buffer...")
+	err := a.recorder.Flush()
+	if err != nil {
+		a.logger.Println(err)
+	}
+}
+
 func generateAgentID() string {
 	agentId, err := uuid.NewRandom()
 	if err != nil {

agent/recorder.go

@@ -203,6 +203,15 @@ func (r *SpanRecorder) Stop() {
 	}
 }
 
+// Flush recorder
+func (r *SpanRecorder) Flush() error {
+	if r.debugMode {
+		r.logger.Println("Flushing recorder buffer...")
+	}
+	err, _ := r.sendSpans()
+	return err
+}
+
 // Write statistics
 func (r *SpanRecorder) writeStats() {
 	r.statsOnce.Do(func() {

init.go

@@ -14,6 +14,7 @@ import (
 	"go.undefinedlabs.com/scopeagent/instrumentation"
 	"go.undefinedlabs.com/scopeagent/instrumentation/logging"
 	scopetesting "go.undefinedlabs.com/scopeagent/instrumentation/testing"
+	"go.undefinedlabs.com/scopeagent/reflection"
 )
 
 var (
@@ -50,6 +51,15 @@ func Run(m *testing.M, opts ...agent.Option) int {
 		newAgent.Stop()
 		os.Exit(1)
 	}()
+	reflection.AddPanicHandler(func(e interface{}) {
+		instrumentation.Logger().Printf("Panic handler triggered by: %v,\nFlushing agent, sending partial results...", e)
+		newAgent.Flush()
+	})
+	reflection.AddOnPanicExitHandler(func(e interface{}) {
+		instrumentation.Logger().Printf("Process is going to end by: %v,\nStopping agent...", e)
+		scopetesting.PanicAllRunningTests(e, 3)
+		newAgent.Stop()
+	})
 
 	defaultAgent = newAgent
 	return newAgent.Run(m)

instrumentation/testing/testing.go

@@ -245,6 +245,29 @@ func GetTest(t *testing.T) *Test {
 	}
 }
 
+// Fails and write panic on running tests
+// Use this only if the process is going to crash
+func PanicAllRunningTests(e interface{}, skip int) {
+	autoInstrumentedTestsMutex.Lock()
+	defer autoInstrumentedTestsMutex.Unlock()
+
+	// We copy the testMap because v.end() locks
+	testMapMutex.RLock()
+	tmp := map[*testing.T]*Test{}
+	for k, v := range testMap {
+		tmp[k] = v
+	}
+	testMapMutex.RUnlock()
+
+	for _, v := range tmp {
+		delete(autoInstrumentedTests, v.t)
+		v.t.Fail()
+		v.span.SetTag("error", true)
+		errors.LogError(v.span, e, 1+skip)
+		v.end()
+	}
+}
+
 // Adds an auto instrumented test to the map
 func addAutoInstrumentedTest(t *testing.T) {
 	autoInstrumentedTestsMutex.Lock()

reflection/panic_handler.go

@@ -0,0 +1,84 @@
+package reflection
+
+import (
+	"reflect"
+	"sync"
+	"unsafe"
+	_ "unsafe"
+
+	"github.com/undefinedlabs/go-mpatch"
+)
+
+var (
+	patchOnPanic    *mpatch.Patch
+	mOnPanic        sync.Mutex
+	onPanicHandlers []func(e interface{})
+
+	patchOnExit   *mpatch.Patch
+	mOnExit       sync.Mutex
+	onExitHandler []func(e interface{})
+)
+
+// Adds a global panic handler (this handler will be executed before any recover call)
+func AddPanicHandler(fn func(interface{})) {
+	mOnPanic.Lock()
+	defer mOnPanic.Unlock()
+	if patchOnPanic == nil {
+		gp := lgopanic
+		np, err := mpatch.PatchMethodByReflect(reflect.Method{Func: reflect.ValueOf(gp)}, gopanic)
+		if err == nil {
+			patchOnPanic = np
+		}
+	}
+	onPanicHandlers = append(onPanicHandlers, fn)
+}
+
+// Adds a global panic handler before process kill (this handler will be executed if not recover is set before the process exits)
+func AddOnPanicExitHandler(fn func(interface{})) {
+	mOnExit.Lock()
+	defer mOnExit.Unlock()
+	if patchOnExit == nil {
+		gp := lpreprintpanics
+		np, err := mpatch.PatchMethodByReflect(reflect.Method{Func: reflect.ValueOf(gp)}, preprintpanics)
+		if err == nil {
+			patchOnExit = np
+		}
+	}
+	onExitHandler = append(onExitHandler, fn)
+}
+
+func gopanic(e interface{}) {
+	mOnPanic.Lock()
+	defer mOnPanic.Unlock()
+	for _, fn := range onPanicHandlers {
+		fn(e)
+	}
+	patchOnPanic.Unpatch()
+	defer patchOnPanic.Patch()
+	lgopanic(e)
+}
+
+func preprintpanics(p *_panic) {
+	mOnExit.Lock()
+	defer mOnExit.Unlock()
+	for _, fn := range onExitHandler {
+		fn(p.arg)
+	}
+	patchOnExit.Unpatch()
+	defer patchOnExit.Patch()
+	lpreprintpanics(p)
+}
+
+//go:linkname lgopanic runtime.gopanic
+func lgopanic(e interface{})
+
+//go:linkname lpreprintpanics runtime.preprintpanics
+func lpreprintpanics(p *_panic)
+
+type _panic struct {
+	argp      unsafe.Pointer // pointer to arguments of deferred call run during panic; cannot move - known to liblink
+	arg       interface{}    // argument to panic
+	link      *_panic        // link to earlier panic
+	recovered bool           // whether this panic is over
+	aborted   bool           // the panic was aborted
+}

reflection/panic_handler.s

@@ -0,0 +1,40 @@
+package reflection_test
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	_ "go.undefinedlabs.com/scopeagent/autoinstrument"
+	"go.undefinedlabs.com/scopeagent/reflection"
+)
+
+func TestPanicHandler(t *testing.T) {
+	var panicHandlerVisit int32
+
+	reflection.AddPanicHandler(func(e interface{}) {
+		t.Log("PANIC HANDLER FOR:", e)
+		atomic.AddInt32(&panicHandlerVisit, 1)
+	})
+
+	t.Run("OnPanic", func(t2 *testing.T) {
+		go func() {
+
+			defer func() {
+				if r := recover(); r != nil {
+					t.Log("PANIC RECOVERED")
+				}
+			}()
+
+			t.Log("PANICKING!")
+			panic("Panic error")
+
+		}()
+
+		time.Sleep(1 * time.Second)
+	})
+
+	if atomic.LoadInt32(&panicHandlerVisit) != 1 {
+		t.Fatalf("panic handler should be executed once.")
+	}
+}