diff --git a/agent/agent.go b/agent/agent.go index b6a0803b..0490058f 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -393,6 +393,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 { diff --git a/agent/recorder.go b/agent/recorder.go index 2cdcbce0..400e085a 100644 --- a/agent/recorder.go +++ b/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() { diff --git a/init.go b/init.go index fb3ae83d..bccaff4d 100644 --- a/init.go +++ b/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) diff --git a/instrumentation/testing/testing.go b/instrumentation/testing/testing.go index 54cc18b7..c72e9167 100644 --- a/instrumentation/testing/testing.go +++ b/instrumentation/testing/testing.go @@ -246,6 +246,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() diff --git a/reflection/panic_handler.go b/reflection/panic_handler.go new file mode 100644 index 00000000..ae575941 --- /dev/null +++ b/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 +} diff --git a/reflection/panic_handler.s b/reflection/panic_handler.s new file mode 100644 index 00000000..e69de29b diff --git a/reflection/panic_handler_test.go b/reflection/panic_handler_test.go new file mode 100644 index 00000000..ec3a0eb1 --- /dev/null +++ b/reflection/panic_handler_test.go @@ -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.") + } +}