main: wait for interrupt to finish before ending command invocation

If a command invocation such as 'daemon' is interrupted, the interrupt
handler asks the node to close. The closing of the node will result in
the command invocation finishing, and possibly returning from main()
before the interrupt handler is done. In particular, the info logging
that a graceful shutdown was completed may never reach reach stdout.

As the whole point of logging "Gracefully shut down." is to give
confidence when debugging that the shutdown was clean, this is
slightly unfortunate.

The interrupt handler needs to be set up in main() instead of Run(),
so that we can defer the closing of the interrupt handler until just
before returning from main, not when Run() returns with a streaming
result reader.

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) {