Merge pull request #60 from m-lab/pt-backup

Add new mode -scamper-with-paris-backup

Author: pboothe

caller.go

@@ -34,7 +34,7 @@ var (
 	waitTime          = flag.Duration("waitTime", 5*time.Second, "how long to wait between subsequent listings of open connections")
 	poll              = flag.Bool("poll", true, "Whether the polling method should be used to see new connections.")
 	tracerType        = flagx.Enum{
-		Options: []string{"paris-traceroute", "scamper"},
+		Options: []string{"paris-traceroute", "scamper", "scamper-with-paris-backup"},
 		Value:   "scamper",
 	}
 
@@ -63,34 +63,46 @@ func main() {
 	promSrv := prometheusx.MustServeMetrics()
 	defer promSrv.Shutdown(ctx)
 
-	var trace ipcache.Tracer
+	scamperDaemon := &tracer.ScamperDaemon{
+		Binary:           *scamperBin,
+		AttachBinary:     *scattachBin,
+		Warts2JSONBinary: *scwarts2jsonBin,
+		OutputPath:       *outputPath,
+		ControlSocket:    *scamperCtrlSocket,
+		ScamperTimeout:   *scamperTimeout,
+	}
+	parisTracer := &tracer.Paris{
+		Binary:     *parisBin,
+		OutputPath: *outputPath,
+		Timeout:    *parisTimeout,
+	}
+
+	var cache *ipcache.RecentIPCache
+
+	// Set up the cache three different ways, depending on the trace method requested.
 	switch tracerType.Value {
 	case "scamper":
-		daemon := &tracer.ScamperDaemon{
-			Binary:           *scamperBin,
-			AttachBinary:     *scattachBin,
-			Warts2JSONBinary: *scwarts2jsonBin,
-			OutputPath:       *outputPath,
-			ControlSocket:    *scamperCtrlSocket,
-			ScamperTimeout:   *scamperTimeout,
-		}
+		cache = ipcache.New(ctx, scamperDaemon, *ipcache.IPCacheTimeout, *ipcache.IPCacheUpdatePeriod)
 		wg.Add(1)
 		go func() {
-			daemon.MustStart(ctx)
+			scamperDaemon.MustStart(ctx)
+			// When the scamper daemon dies, cancel main() and exit.
 			cancel()
 			wg.Done()
 		}()
-		trace = daemon
 	case "paris-traceroute":
-		trace = &tracer.Paris{
-			Binary:     *parisBin,
-			OutputPath: *outputPath,
-			Timeout:    *parisTimeout,
-		}
+		cache = ipcache.New(ctx, parisTracer, *ipcache.IPCacheTimeout, *ipcache.IPCacheUpdatePeriod)
+	case "scamper-with-paris-backup":
+		cache = ipcache.New(ctx, scamperDaemon, *ipcache.IPCacheTimeout, *ipcache.IPCacheUpdatePeriod)
+		wg.Add(1)
+		go func() {
+			scamperDaemon.MustStart(ctx)
+			// When the scamper daemon dies, switch to paris-traceroute.
+			cache.UpdateTracer(parisTracer)
+			wg.Done()
+		}()
 	}
 
-	cache := ipcache.New(ctx, trace, *ipcache.IPCacheTimeout, *ipcache.IPCacheUpdatePeriod)
-
 	if *poll {
 		wg.Add(1)
 		go func(c *ipcache.RecentIPCache) {

caller_test.go

@@ -32,6 +32,7 @@ func TestMain(t *testing.T) {
 	*waitTime = time.Nanosecond // Run through the loop a few times.
 	*outputPath = dir
 	*poll = true
+	*scamperBin = "scamper"
 	tracerType.Value = "scamper"
 	ctx, cancel = context.WithCancel(context.Background())
 	go func() {
@@ -49,7 +50,6 @@ func TestMainWithConnectionListener(t *testing.T) {
 	rtx.Must(srv.Listen(), "Could not start the empty server")
 
 	*prometheusx.ListenAddress = ":0"
-	*scamperCtrlSocket = dir + "/scamper.sock"
 	*eventsocket.Filename = dir + "/events.sock"
 	*outputPath = dir
 	*poll = false
@@ -64,6 +64,30 @@ func TestMainWithConnectionListener(t *testing.T) {
 	main()
 }
 
+func TestMainWithBackupPT(t *testing.T) {
+	dir, err := ioutil.TempDir("", "TestMainWithBackupPT")
+	rtx.Must(err, "Could not create temp dir")
+	defer os.RemoveAll(dir)
+	srv := eventsocket.New(dir + "/events.sock")
+	rtx.Must(srv.Listen(), "Could not start the empty server")
+
+	*prometheusx.ListenAddress = ":0"
+	*eventsocket.Filename = dir + "/events.sock"
+	*outputPath = dir
+	*poll = false
+	*scamperCtrlSocket = dir + "/scamper.sock"
+	*scamperBin = "false"
+	tracerType.Value = "scamper-with-paris-backup"
+
+	ctx, cancel = context.WithCancel(context.Background())
+	go srv.Serve(ctx)
+	go func() {
+		time.Sleep(1 * time.Second)
+		cancel()
+	}()
+	main()
+}
+
 func TestMainWithBadArgs(t *testing.T) {
 	tracerType.Value = "paris-traceroute"
 	*eventsocket.Filename = ""

ipcache/ipcache.go

@@ -56,21 +56,28 @@ func (rc *RecentIPCache) getEntry(ip string) (*cachedTest, bool) {
 	return rc.cache[ip], existed
 }
 
+func (rc *RecentIPCache) getTracer() Tracer {
+	rc.mu.Lock()
+	defer rc.mu.Unlock()
+	return rc.tracer
+}
+
 // Trace performs a trace and adds it to the cache. It calls the methods of the
 // tracer, so if those create files on disk, then files on disk will be created
 // as a side effect.
 func (rc *RecentIPCache) Trace(conn connection.Connection) (string, error) {
 	c, cached := rc.getEntry(conn.RemoteIP)
+	t := rc.getTracer()
 	if cached {
 		<-c.dataReady
 		if c.err == nil {
-			rc.tracer.TraceFromCachedTrace(conn, time.Now(), c.data)
+			t.TraceFromCachedTrace(conn, time.Now(), c.data)
 			return c.data, nil
 		}
-		rc.tracer.DontTrace(conn, c.err)
+		t.DontTrace(conn, c.err)
 		return "", c.err
 	}
-	c.data, c.err = rc.tracer.Trace(conn, c.timeStamp)
+	c.data, c.err = t.Trace(conn, c.timeStamp)
 	close(c.dataReady)
 	return c.data, c.err
 }
@@ -84,6 +91,14 @@ func (rc *RecentIPCache) GetCacheLength() int {
 	return len(rc.cache)
 }
 
+// UpdateTracer switches the Tracer being used. This allows us to dynamically
+// switch between scamper and paris-traceroute.
+func (rc *RecentIPCache) UpdateTracer(t Tracer) {
+	rc.mu.Lock()
+	defer rc.mu.Unlock()
+	rc.tracer = t
+}
+
 // New creates and returns a RecentIPCache. It also starts up a background
 // goroutine that scrubs the cache.
 func New(ctx context.Context, trace Tracer, ipCacheTimeout, ipCacheUpdatePeriod time.Duration) *RecentIPCache {

ipcache/ipcache_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/m-lab/go/rtx"
 	"github.com/m-lab/traceroute-caller/connection"
 	"github.com/m-lab/traceroute-caller/ipcache"
 	pipe "gopkg.in/m-lab/pipe.v3"
@@ -23,6 +24,7 @@ type testTracer struct {
 }
 
 func (tf *testTracer) Trace(conn connection.Connection, t time.Time) (string, error) {
+	tf.calls++
 	return "Fake trace test " + conn.RemoteIP, nil
 }
 
@@ -82,6 +84,35 @@ func TestTrace(t *testing.T) {
 	}
 }
 
+func TestUpdateTracer(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	var tt, tt2 testTracer
+	testCache := ipcache.New(ctx, &tt, 100*time.Second, time.Second)
+	conn1 := connection.Connection{
+		RemoteIP:   "1.1.1.2",
+		RemotePort: 5034,
+		LocalIP:    "1.1.1.3",
+		LocalPort:  58790,
+		Cookie:     "10f3d"}
+	_, err := testCache.Trace(conn1)
+	rtx.Must(err, "Could not trace using tt")
+
+	testCache.UpdateTracer(&tt2)
+	conn2 := connection.Connection{
+		RemoteIP:   "1.1.1.5",
+		RemotePort: 5034,
+		LocalIP:    "1.1.1.3",
+		LocalPort:  58790,
+		Cookie:     "aaaa"}
+	_, err = testCache.Trace(conn2)
+	rtx.Must(err, "Could not trace using tt2")
+
+	if tt.calls != 1 || tt2.calls != 1 {
+		t.Error("Each tracer should have been called once, not", tt.calls, tt2.calls)
+	}
+}
+
 func TestRecentIPCache(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()