From 7bbbab87742f79130bee20379c165cdaba03a13a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arturo=20Filast=C3=B2?= <arturo@openobservatory.org>
Date: Fri, 27 Dec 2019 11:32:08 +0100
Subject: [PATCH] Handle the SIGINT and SIGTERM signals to support stopping a
 test cleanly (#84)

---
 internal/cli/run/run.go | 25 ++++++++++++++++++++++++-
 nettests/nettests.go    |  5 ++++-
 ooni.go                 | 23 ++++++++++++++++++++---
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/internal/cli/run/run.go b/internal/cli/run/run.go
index 066c205..0116db8 100644
--- a/internal/cli/run/run.go
+++ b/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
 		}
 	}
 
diff --git a/nettests/nettests.go b/nettests/nettests.go
index a5f3a72..4826598 100644
--- a/nettests/nettests.go
+++ b/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"))
diff --git a/ooni.go b/ooni.go
index d799222..8f6ac90 100644
--- a/ooni.go
+++ b/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
@@ -94,9 +110,10 @@ func (c *Context) Init() error {
 // NewContext creates a new context instance.
 func NewContext(configPath string, homePath string) *Context {
 	return &Context{
-		Home:       homePath,
-		Config:     &config.Config{},
-		configPath: configPath,
+		Home:                  homePath,
+		Config:                &config.Config{},
+		configPath:            configPath,
+		isTerminatedAtomicInt: 0,
 	}
 }