Handle the SIGINT and SIGTERM signals to support stopping a test cleanly (#84)

internal/cli/run/run.go

@@ -2,6 +2,9 @@ package run
 
 import (
 	"errors"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/alecthomas/kingpin"
 	"github.com/apex/log"
@@ -13,6 +16,22 @@ import (
 	"github.com/ooni/probe-cli/nettests"
 )
 
+// listenForSignals will listen for SIGINT and SIGTERM. When it receives those
+// signals it will set isTerminatedAtomicInt to non-zero, which will cleanly
+// shutdown the test logic.
+// TODO refactor this to use a cancellable context.Context instead of a bool
+// flag, probably as part of: https://github.com/ooni/probe-cli/issues/45
+func listenForSignals(ctx *ooni.Context) {
+	s := make(chan os.Signal, 1)
+	signal.Notify(s, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		<-s
+		log.Info("caught a stop signal, shutting down cleanly")
+		ctx.Terminate()
+	}()
+}
+
 func runNettestGroup(tg string, ctx *ooni.Context, network *database.Network) error {
 	group, ok := nettests.NettestGroups[tg]
 	if !ok {
@@ -27,13 +46,17 @@ func runNettestGroup(tg string, ctx *ooni.Context, network *database.Network) er
 		return err
 	}
 
+	listenForSignals(ctx)
 	for i, nt := range group.Nettests {
+		if ctx.IsTerminated() == true {
+			log.Debugf("context is terminated, breaking")
+			break
+		}
 		log.Debugf("Running test %T", nt)
 		ctl := nettests.NewController(nt, ctx, result)
 		ctl.SetNettestIndex(i, len(group.Nettests))
 		if err = nt.Run(ctl); err != nil {
 			log.WithError(err).Errorf("Failed to run %s", group.Label)
-			return err
 		}
 	}
 

nettests/nettests.go

@@ -77,7 +77,6 @@ func (c *Controller) SetNettestIndex(i, n int) {
 // This function will continue to run in most cases but will
 // immediately halt if something's wrong with the file system.
 func (c *Controller) Run(builder *engine.ExperimentBuilder, inputs []string) error {
-
 	// This will configure the controller as handler for the callbacks
 	// called by ooni/probe-engine/experiment.Experiment.
 	builder.SetCallbacks(engine.Callbacks(c))
@@ -108,6 +107,10 @@ func (c *Controller) Run(builder *engine.ExperimentBuilder, inputs []string) err
 
 	c.ntStartTime = time.Now()
 	for idx, input := range inputs {
+		if c.Ctx.IsTerminated() == true {
+			log.Debug("isTerminated == true, breaking the input loop")
+			break
+		}
 		c.curInputIdx = idx // allow for precise progress
 		idx64 := int64(idx)
 		log.Debug(color.RedString("status.measurement_start"))

ooni.go

@@ -3,6 +3,7 @@ package ooni
 import (
 	"io/ioutil"
 	"os"
+	"sync/atomic"
 
 	"github.com/apex/log"
 	"github.com/ooni/probe-cli/config"
@@ -29,6 +30,10 @@ type Context struct {
 
 	dbPath     string
 	configPath string
+
+	// We need to use a int64 in order to use the atomic.AddInt64/LoadInt64
+	// operations to ensure consistent reads of the variables.
+	isTerminatedAtomicInt int64
 }
 
 // MaybeLocationLookup will lookup the location of the user unless it's already cached
@@ -36,6 +41,17 @@ func (c *Context) MaybeLocationLookup() error {
 	return c.Session.MaybeLookupLocation()
 }
 
+// IsTerminated checks to see if the isTerminatedAtomicInt is set to a non zero
+// value and therefore we have received the signal to shutdown the running test
+func (c *Context) IsTerminated() bool {
+	i := atomic.LoadInt64(&c.isTerminatedAtomicInt)
+	return i != 0
+}
+
+func (c *Context) Terminate() {
+	atomic.AddInt64(&c.isTerminatedAtomicInt, 1)
+}
+
 // Init the OONI manager
 func (c *Context) Init() error {
 	var err error
@@ -97,6 +113,7 @@ func NewContext(configPath string, homePath string) *Context {
 		Home:                  homePath,
 		Config:                &config.Config{},
 		configPath:            configPath,
+		isTerminatedAtomicInt: 0,
 	}
 }