TOOLS-1341 consistent signal handling

common/signals/signals.go

@@ -0,0 +1,51 @@
+package signals
+
+import (
+	"github.com/mongodb/mongo-tools/common/log"
+	"github.com/mongodb/mongo-tools/common/util"
+
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// Handle is like HandleWithInterrupt but it doesn't take a finalizer and will
+// exit immediately after the first signal is received.
+func Handle() chan struct{} {
+	return HandleWithInterrupt(nil)
+}
+
+// HandleWithInterrupt starts a goroutine which listens for SIGTERM, SIGINT,
+// SIGKILL, and SIGPIPE. It calls the finalizer function when the first signal
+// is received and forcibly terminates the program after the second. If a nil
+// function is provided, the program will exit after the first signal.
+func HandleWithInterrupt(finalizer func()) chan struct{} {
+	finishedChan := make(chan struct{})
+	go handleSignals(finalizer, finishedChan)
+	return finishedChan
+}
+
+func handleSignals(finalizer func(), finishedChan chan struct{}) {
+	log.Logv(log.DebugLow, "will listen for SIGTERM, SIGINT, SIGKILL, and SIGPIPE")
+	sigChan := make(chan os.Signal, 2)
+	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL, syscall.SIGPIPE)
+	defer signal.Stop(sigChan)
+	if finalizer != nil {
+		select {
+		case <-sigChan:
+			// first signal use finalizer to terminate cleanly
+			log.Logv(log.Always, "signal received; attempting to shut down")
+			finalizer()
+		case <-finishedChan:
+			return
+		}
+	}
+	select {
+	case <-sigChan:
+		// second signal exits immediately
+		log.Logv(log.Always, "signal received; forcefully terminating")
+		os.Exit(util.ExitKill)
+	case <-finishedChan:
+		return
+	}
+}

mongodump/main/mongodump.go

@@ -2,11 +2,13 @@
 package main
 
 import (
+	"os"
+
 	"github.com/mongodb/mongo-tools/common/log"
 	"github.com/mongodb/mongo-tools/common/options"
+	"github.com/mongodb/mongo-tools/common/signals"
 	"github.com/mongodb/mongo-tools/common/util"
 	"github.com/mongodb/mongo-tools/mongodump"
-	"os"
 )
 
 func main() {
@@ -55,6 +57,9 @@ func main() {
 		InputOptions:  inputOpts,
 	}
 
+	finishedChan := signals.HandleWithInterrupt(dump.HandleInterrupt)
+	defer close(finishedChan)
+
 	if err = dump.Init(); err != nil {
 		log.Logvf(log.Always, "Failed: %v", err)
 		os.Exit(util.ExitError)

mongodump/mongodump.go

@@ -2,8 +2,6 @@
 package mongodump
 
 import (
-	"compress/gzip"
-	"fmt"
 	"github.com/mongodb/mongo-tools/common/archive"
 	"github.com/mongodb/mongo-tools/common/auth"
 	"github.com/mongodb/mongo-tools/common/bsonutil"
@@ -16,12 +14,13 @@ import (
 	"github.com/mongodb/mongo-tools/common/util"
 	"gopkg.in/mgo.v2"
 	"gopkg.in/mgo.v2/bson"
+
+	"compress/gzip"
+	"fmt"
 	"io"
 	"os"
-	"os/signal"
 	"path/filepath"
 	"sync"
-	"syscall"
 	"time"
 )
 
@@ -175,7 +174,6 @@ func (dump *MongoDump) Init() error {
 
 // Dump handles some final options checking and executes MongoDump.
 func (dump *MongoDump) Dump() (err error) {
-
 	dump.shutdownIntentsNotifier = newNotifier()
 
 	if dump.InputOptions.HasQuery() {
@@ -375,8 +373,6 @@ func (dump *MongoDump) Dump() (err error) {
 	dump.progressManager.Start()
 	defer dump.progressManager.Stop()
 
-	go dump.handleSignals()
-
 	if err := dump.DumpIntents(); err != nil {
 		return err
 	}
@@ -794,26 +790,14 @@ func (dump *MongoDump) getArchiveOut() (out io.WriteCloser, err error) {
 	return out, nil
 }
 
-// handleSignals listens for either SIGTERM, SIGINT or the
-// SIGHUP signal. It ends restore reads for all goroutines
-// as soon as any of those signals is received.
-func (dump *MongoDump) handleSignals() {
-	log.Logv(log.DebugLow, "will listen for SIGTERM, SIGINT and SIGHUP")
-	sigChan := make(chan os.Signal, 2)
-	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP, syscall.SIGPIPE)
-	// first signal cleanly terminates dump writes
-	<-sigChan
-	log.Logv(log.Always, "signal received; shutting down mongodump")
-	//
-	dump.shutdownIntentsNotifier.Notify()
-	// second signal exits immediately
-	<-sigChan
-	log.Logv(log.Always, "second signal received; forcefully terminating mongodump")
-	os.Exit(util.ExitKill)
-}
-
 // docPlural returns "document" or "documents" depending on the
 // count of documents passed in.
 func docPlural(count int64) string {
 	return util.Pluralize(int(count), "document", "documents")
 }
+
+func (dump *MongoDump) HandleInterrupt() {
+	if dump.shutdownIntentsNotifier != nil {
+		dump.shutdownIntentsNotifier.Notify()
+	}
+}

mongoexport/main/mongoexport.go

@@ -14,7 +14,7 @@ import (
 )
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 	// initialize command-line opts
 	opts := options.New("mongoexport", mongoexport.Usage,
 		options.EnabledOptions{Auth: true, Connection: true, Namespace: true})

mongofiles/main/mongofiles.go

@@ -13,7 +13,7 @@ import (
 )
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 	// initialize command-line opts
 	opts := options.New("mongofiles", mongofiles.Usage, options.EnabledOptions{Auth: true, Connection: true, Namespace: false})
 

mongoimport/main/mongoimport.go

@@ -14,7 +14,7 @@ import (
 )
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 	// initialize command-line opts
 	opts := options.New("mongoimport", mongoimport.Usage,
 		options.EnabledOptions{Auth: true, Connection: true, Namespace: true})

mongooplog/main/mongooplog.go

@@ -12,7 +12,7 @@ import (
 )
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 
 	// initialize command line options
 	opts := options.New("mongooplog", mongooplog.Usage,

mongorestore/main/mongorestore.go

@@ -8,6 +8,7 @@ import (
 	"github.com/mongodb/mongo-tools/common/db"
 	"github.com/mongodb/mongo-tools/common/log"
 	"github.com/mongodb/mongo-tools/common/options"
+	"github.com/mongodb/mongo-tools/common/signals"
 	"github.com/mongodb/mongo-tools/common/util"
 	"github.com/mongodb/mongo-tools/mongorestore"
 )
@@ -76,6 +77,9 @@ func main() {
 		SessionProvider: provider,
 	}
 
+	finishedChan := signals.HandleWithInterrupt(restore.HandleInterrupt)
+	defer close(finishedChan)
+
 	if err = restore.Restore(); err != nil {
 		log.Logvf(log.Always, "Failed: %v", err)
 		if err == util.ErrTerminated {

mongorestore/mongorestore.go

@@ -7,10 +7,8 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
-	"os/signal"
 	"path/filepath"
 	"sync"
-	"syscall"
 
 	"github.com/mongodb/mongo-tools/common/archive"
 	"github.com/mongodb/mongo-tools/common/auth"
@@ -462,7 +460,6 @@ func (restore *MongoRestore) Restore() error {
 	}
 
 	restore.termChan = make(chan struct{})
-	go restore.handleSignals()
 
 	if err := restore.RestoreIntents(); err != nil {
 		return err
@@ -485,6 +482,7 @@ func (restore *MongoRestore) Restore() error {
 	}
 
 	log.Logv(log.Always, "done")
+
 	return nil
 }
 
@@ -535,19 +533,8 @@ func (restore *MongoRestore) getArchiveReader() (rc io.ReadCloser, err error) {
 	return rc, nil
 }
 
-// handleSignals listens for either SIGTERM, SIGINT or the
-// SIGHUP signal. It ends restore reads for all goroutines
-// as soon as any of those signals is received.
-func (restore *MongoRestore) handleSignals() {
-	log.Logv(log.DebugLow, "will listen for SIGTERM and SIGINT")
-	sigChan := make(chan os.Signal, 2)
-	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
-	// first signal cleanly terminates restore reads
-	<-sigChan
-	log.Logv(log.Always, "ending restore reads")
-	close(restore.termChan)
-	// second signal exits immediately
-	<-sigChan
-	log.Logv(log.Always, "forcefully terminating mongorestore")
-	os.Exit(util.ExitKill)
+func (restore *MongoRestore) HandleInterrupt() {
+	if restore.termChan != nil {
+		close(restore.termChan)
+	}
 }

mongostat/main/mongostat.go

@@ -47,7 +47,7 @@ func optionCustomHeaders(option string) (headers []string) {
 }
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 	// initialize command-line opts
 	opts := options.New(
 		"mongostat",

mongotop/main/mongotop.go

@@ -15,7 +15,7 @@ import (
 )
 
 func main() {
-	go signals.Handle()
+	signals.Handle()
 
 	// initialize command-line opts
 	opts := options.New("mongotop", mongotop.Usage,