diff --git a/cmd/ipfs/main.go b/cmd/ipfs/main.go index 5cee9f8c7e4c4dd4f0d4b7215cb923eaf6b4a21f..1ba5f662b6876c466814c05facb61624248d1c30 100644 --- a/cmd/ipfs/main.go +++ b/cmd/ipfs/main.go @@ -11,6 +11,7 @@ import ( "runtime" "runtime/pprof" "strings" + "sync" "syscall" "time" @@ -39,7 +40,6 @@ const ( cpuProfile = "ipfs.cpuprof" heapProfile = "ipfs.memprof" errorFormat = "ERROR: %v\n\n" - shutdownMessage = "Received interrupt signal, shutting down..." ) type cmdInvocation struct { @@ -141,6 +141,8 @@ func main() { } // ok, finally, run the command invocation. + intrh := invoc.SetupInterruptHandler() + defer intrh.Close() output, err := invoc.Run(ctx) if err != nil { printErr(err) @@ -157,8 +159,6 @@ func main() { } func (i *cmdInvocation) Run(ctx context.Context) (output io.Reader, err error) { - // setup our global interrupt handler. - i.setupInterruptHandler() // check if user wants to debug. option OR env var. debug, _, err := i.req.Option("debug").Bool() @@ -474,57 +474,87 @@ func writeHeapProfileToFile() error { return pprof.WriteHeapProfile(mprof) } -// listen for and handle SIGTERM -func (i *cmdInvocation) setupInterruptHandler() { +// IntrHandler helps set up an interrupt handler that can +// be cleanly shut down through the io.Closer interface. +type IntrHandler struct { + sig chan os.Signal + wg sync.WaitGroup +} + +func NewIntrHandler() *IntrHandler { + ih := &IntrHandler{} + ih.sig = make(chan os.Signal, 1) + return ih +} + +func (ih *IntrHandler) Close() error { + close(ih.sig) + ih.wg.Wait() + return nil +} - ctx := i.req.Context() - sig := allInterruptSignals() +// Handle starts handling the given signals, and will call the handler +// callback function each time a signal is catched. The function is passed +// the number of times the handler has been triggered in total, as +// well as the handler itself, so that the handling logic can use the +// handler's wait group to ensure clean shutdown when Close() is called. +func (ih *IntrHandler) Handle(handler func(count int, ih *IntrHandler), sigs ...os.Signal) { + signal.Notify(ih.sig, sigs...) + ih.wg.Add(1) go func() { - // first time, try to shut down. + defer ih.wg.Done() + count := 0 + for _ = range ih.sig { + count++ + handler(count, ih) + } + signal.Stop(ih.sig) + }() +} + +func (i *cmdInvocation) SetupInterruptHandler() io.Closer { - // loop because we may be - for count := 0; ; count++ { - <-sig + intrh := NewIntrHandler() + handlerFunc := func(count int, ih *IntrHandler) { + switch count { + case 1: + // first time, try to shut down + fmt.Println("Received interrupt signal, shutting down...") + + ctx := i.req.Context() // if we're still initializing, cannot use `ctx.GetNode()` select { default: // initialization not done - fmt.Println(shutdownMessage) os.Exit(-1) case <-ctx.InitDone: } - // TODO cancel the command context instead + ih.wg.Add(1) + go func() { + defer ih.wg.Done() - n, err := ctx.GetNode() - if err != nil { - log.Error(err) - fmt.Println(shutdownMessage) - os.Exit(-1) - } + // TODO cancel the command context instead + n, err := ctx.GetNode() + if err != nil { + log.Error(err) + os.Exit(-1) + } - switch count { - case 0: - fmt.Println(shutdownMessage) - go func() { - n.Close() - log.Info("Gracefully shut down.") - }() + n.Close() + log.Info("Gracefully shut down.") + }() - default: - fmt.Println("Received another interrupt before graceful shutdown, terminating...") - os.Exit(-1) - } + default: + fmt.Println("Received another interrupt before graceful shutdown, terminating...") + os.Exit(-1) } - }() -} + } + + intrh.Handle(handlerFunc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM) -func allInterruptSignals() chan os.Signal { - sigc := make(chan os.Signal, 1) - signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT, - syscall.SIGTERM) - return sigc + return intrh } func profileIfEnabled() (func(), error) {