Skip to content

Commit

Permalink
Add signals.RegisterStackTraceHandlerOnSignals for custom handling lo…
Browse files Browse the repository at this point in the history
…gic (#129)
  • Loading branch information
bmoylan authored Dec 11, 2018
1 parent b99e4c1 commit 97c8883
Showing 1 changed file with 21 additions and 7 deletions.
28 changes: 21 additions & 7 deletions signals/signals.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
package signals

import (
"bytes"
"context"
"io"
"os"
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 97c8883

Please sign in to comment.