From 97c88835cd7ba5e284c4d94f1b63e0709a78bc01 Mon Sep 17 00:00:00 2001
From: Brad Moylan <moylan.brad@gmail.com>
Date: Mon, 10 Dec 2018 16:43:53 -0800
Subject: [PATCH] Add signals.RegisterStackTraceHandlerOnSignals for custom
 handling logic (#129)

---
 signals/signals.go | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/signals/signals.go b/signals/signals.go
index 4715eeb4..0d441631 100644
--- a/signals/signals.go
+++ b/signals/signals.go
@@ -5,6 +5,7 @@
 package signals
 
 import (
+	"bytes"
 	"context"
 	"io"
 	"os"
@@ -60,26 +61,39 @@ func RegisterStackTraceWriter(out io.Writer, errHandler func(error)) (unregister
 // error, that error is provided to the errHandler function if one is provided. Returns a function that unregisters the
 // listener when called.
 func RegisterStackTraceWriterOnSignals(out io.Writer, errHandler func(error), sig ...os.Signal) (unregister func()) {
-	cancel := make(chan bool, 1)
-	unregister = func() {
-		cancel <- true
+	ctx, cancel := context.WithCancel(context.Background())
+	handler := func(stackTraceOutput []byte) error {
+		_, err := out.Write(stackTraceOutput)
+		return err
 	}
+	RegisterStackTraceHandlerOnSignals(ctx, handler, errHandler, sig...)
+	return cancel
+}
 
+// RegisterStackTraceHandlerOnSignals starts a goroutine that listens for the specified signals and calls stackTraceHandler with a
+// pprof-formatted snapshot of all running goroutines when any of the provided signals are received. If stackTraceHandler returns
+// an error, that error is provided to the errHandler function if one is provided. No goroutine is created if stackTraceHandler is nil.
+// If no signals are provided, the handler will receive notifications for all signals (matching the os/signal.Notify API).
+// The handler will exit when ctx is cancelled.
+func RegisterStackTraceHandlerOnSignals(ctx context.Context, stackTraceHandler func(stackTraceOutput []byte) error, errHandler func(error), sig ...os.Signal) {
+	if stackTraceHandler == nil {
+		return
+	}
 	signals := NewSignalReceiver(sig...)
 	go func() {
 		for {
 			select {
 			case <-signals:
-				err := pprof.Lookup("goroutine").WriteTo(out, 2)
-				if err != nil && errHandler != nil {
+				var buf bytes.Buffer
+				_ = pprof.Lookup("goroutine").WriteTo(&buf, 2) // bytes.Buffer's Write never returns an error, so we swallow it
+				if err := stackTraceHandler(buf.Bytes()); err != nil && errHandler != nil {
 					errHandler(err)
 				}
-			case <-cancel:
+			case <-ctx.Done():
 				return
 			}
 		}
 	}()
-	return unregister
 }
 
 // NewSignalReceiver returns a buffered channel that is registered to receive the provided signals.