Feature: disconnect agent from API on signal

If the agent receives a SIGUSR1 signal, it will disconnect from the API,
but it will continue running checks. This makes it possible for a new
version of the agent to connect to the API. When that happens, it's
possible to kill the running agent (SIGTERM or otherwise).

This should provide for a smoother upgrade process.

Signed-off-by: Marcelo E. Magallon <[email protected]>

Author: mem

README.md

@@ -84,3 +84,18 @@ Now you should have the agent reporting as private probe, and running checks (if
 
 #### Deploy it using Kubernetes
 See [examples/kubernetes](./examples/kubernetes) for the documentation and example yaml files
+
+Signals
+-------
+
+The agent traps the following signals:
+
+* SIGTERM: The agent tries to clean up and shut down in an orderly
+  manner.
+* SIGUSR1: The agent disconnects from the API but keeps running checks.
+  After 1 minute elapses, the agent will try to reconnect to the API and
+  keep trying until it succeeds or it's killed. One possible use case is
+  upgrading a running agent with a newer version: after SIGUSR1 is sent,
+  the agent disconnects, allowing another agent to connect in its place.
+  If the new agent fails to connect, the old agent will reconnect and
+  take it from there.

internal/checks/checks.go

@@ -5,8 +5,12 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
+	"os/signal"
 	"strconv"
 	"sync"
+	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/feature"
@@ -24,8 +28,9 @@ import (
 )
 
 var (
-	errNotAuthorized    = errors.New("probe not authorized")
-	errTransportClosing = errors.New("transport closing")
+	errNotAuthorized     = errors.New("probe not authorized")
+	errTransportClosing  = errors.New("transport closing")
+	errProbeUnregistered = errors.New("probe no longer registered")
 )
 
 // Updater represents a probe along with the collection of scrapers
@@ -219,16 +224,29 @@ func (c *Updater) Run(ctx context.Context) error {
 				Msg("context cancelled, closing updater")
 			return nil
 
+		case errors.Is(err, errProbeUnregistered):
+			// Probe unregistered itself from the API, wait
+			// until attempting to reconnect again.
+			c.logger.Warn().
+				Str("connection_state", c.api.conn.GetState().String()).
+				Msg("unregistered probe in API, sleeping for 1 minute...")
+
+			if err := sleepCtx(ctx, 1*time.Minute); err != nil {
+				return err
+			}
+
 		default:
 			c.logger.Warn().
 				Err(err).
 				Str("connection_state", c.api.conn.GetState().String()).
 				Msg("handling check changes")
+
 			// TODO(mem): this might be a transient error (e.g. bad connection). We probably need to
 			// fine-tune GRPPC's backoff parameters. We might also need to keep count of the reconnects, and
 			// give up if they hit some threshold?
-			time.Sleep(2 * time.Second)
-			continue
+			if err := sleepCtx(ctx, 2*time.Second); err != nil {
+				return err
+			}
 		}
 	}
 }
@@ -240,7 +258,7 @@ func (c *Updater) loop(ctx context.Context) error {
 
 	grpcErrorHandler := func(action string, err error) error {
 		status, ok := status.FromError(err)
-		c.logger.Error().Err(err).Uint32("code", uint32(status.Code())).Msg(status.Message())
+		c.logger.Error().Err(err).Str("action", action).Uint32("code", uint32(status.Code())).Msg(status.Message())
 
 		switch {
 		case !ok:
@@ -297,18 +315,79 @@ func (c *Updater) loop(ctx context.Context) error {
 		"buildstamp": version.Buildstamp(),
 	}).Set(1)
 
-	cc, err := client.GetChanges(ctx, &sm.Void{})
-	if err != nil {
-		return grpcErrorHandler("requesting changes from synthetic-monitoring-api", err)
+	// Create a child context so that we can communicate to the
+	// signal handler that we are done.
+	sigCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	// We get _another_ context from the signal handler that we can
+	// use tell the GRPC client that we need to break out. We have
+	// multiple ways of cancelling the context (another signal
+	// elsewhere in the system communicated through the parent
+	// context; cancelling the child context because we are
+	// returning from this function; cancelling the new context
+	// because the signal fired), so we need an additional way of
+	// telling them apart.
+	sigCtx, signalFired := installSignalHandler(sigCtx)
+
+	action := "requesting changes from synthetic-monitoring-api"
+	cc, err := client.GetChanges(sigCtx, &sm.Void{})
+	if err == nil {
+		action = "getting changes from synthetic-monitoring-api"
+		// processChanges uses the context in its first
+		// argument to create scrapers. This means that
+		// cancelling that context cancels all the running
+		// scrapers. That's why we are passing the _original_
+		// context, not sigCtx, so that scrapers are _not_
+		// stopped if the signal is trapped. We want scrapers to
+		// continue running in case the agent is _not_ killed.
+		err = c.processChanges(ctx, cc)
 	}
 
-	if err := c.processChanges(ctx, cc); err != nil {
-		return grpcErrorHandler("getting changes from synthetic-monitoring-api", err)
+	if err != nil {
+		if atomic.LoadInt32(signalFired) == 1 {
+			return errProbeUnregistered
+		}
+
+		return grpcErrorHandler(action, err)
 	}
 
 	return nil
 }
 
+// installSignalHandler installs a signal handler for SIGUSR1.
+//
+// The returned context's Done channel is closed if the signal is
+// delivered. To make it simpler to determine if the signal was
+// delivered, a value of 1 is written to the location pointed to by the
+// returned int32 pointer.
+//
+// If the provided context's Done channel is closed before the signal is
+// delivered, the signal handler is removed and the returned context's
+// Done channel is closed, too. It's the callers responsibility to
+// cancel the provided context if it's no longer interested in the
+// signal.
+func installSignalHandler(ctx context.Context) (context.Context, *int32) {
+	sigCtx, cancel := context.WithCancel(ctx)
+
+	fired := new(int32)
+
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGUSR1)
+
+	go func() {
+		select {
+		case <-sigCh:
+			atomic.StoreInt32(fired, 1)
+			cancel()
+		case <-ctx.Done():
+		}
+		signal.Stop(sigCh)
+	}()
+
+	return sigCtx, fired
+}
+
 func (c *Updater) processChanges(ctx context.Context, cc sm.Checks_GetChangesClient) error {
 	firstBatch := true
 
@@ -317,6 +396,9 @@ func (c *Updater) processChanges(ctx context.Context, cc sm.Checks_GetChangesCli
 		case <-cc.Context().Done():
 			return nil
 
+		case <-ctx.Done():
+			return ctx.Err()
+
 		default:
 			switch msg, err := cc.Recv(); err {
 			case nil:
@@ -472,6 +554,10 @@ func (c *Updater) handleFirstBatch(ctx context.Context, changes *sm.Changes) {
 			continue
 		}
 
+		c.logger.Debug().
+			Int64("check_id", id).
+			Msg("stopping scraper during first batch handling")
+
 		checkType := scraper.CheckType().String()
 		scraper.Stop()
 
@@ -599,3 +685,24 @@ func (c *Updater) addAndStartScraperWithLock(ctx context.Context, check sm.Check
 
 	return nil
 }
+
+// sleepCtx is like time.Sleep, but it pays attention to the
+// cancellation of the provided context.
+func sleepCtx(ctx context.Context, d time.Duration) error {
+	var err error
+
+	timer := time.NewTimer(d)
+
+	select {
+	case <-ctx.Done():
+		err = ctx.Err()
+
+		if !timer.Stop() {
+			<-timer.C
+		}
+
+	case <-timer.C:
+	}
+
+	return err
+}

internal/checks/checks_test.go

@@ -1,7 +1,11 @@
 package checks
 
 import (
+	"context"
+	"sync/atomic"
+	"syscall"
 	"testing"
+	"time"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/feature"
 	"github.com/grafana/synthetic-monitoring-agent/internal/pusher"
@@ -53,3 +57,87 @@ func TestNewUpdater(t *testing.T) {
 		})
 	}
 }
+
+func TestInstallSignalHandler(t *testing.T) {
+	testcases := map[string]func(t *testing.T){
+		"signal": func(t *testing.T) {
+			// verify that the signal context is done after
+			// receiving the signal, and that the signal is
+			// correctly reported as having fired.
+
+			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+			defer cancel()
+			sigCtx, signalFired := installSignalHandler(ctx)
+			require.NotNil(t, sigCtx)
+			require.NotNil(t, signalFired)
+			require.NoError(t, syscall.Kill(syscall.Getpid(), syscall.SIGUSR1))
+
+			select {
+			case <-ctx.Done():
+				t.Fatal("context timeout expired")
+			case <-sigCtx.Done():
+				require.Equal(t, int32(1), atomic.LoadInt32(signalFired))
+			}
+		},
+
+		"no signal": func(t *testing.T) {
+			// verify that the signal context is done after
+			// the parrent context is done, and that the
+			// signal is correctly reported as not having
+			// fired.
+
+			ctx, cancel := context.WithCancel(context.Background())
+			sigCtx, signalFired := installSignalHandler(ctx)
+			require.NotNil(t, sigCtx)
+			require.NotNil(t, signalFired)
+
+			cancel()
+
+			timeout := 100 * time.Millisecond
+			timer := time.NewTimer(timeout)
+			defer timer.Stop()
+
+			select {
+			case <-timer.C:
+				t.Fatalf("signal context not cancelled after %s", timeout)
+			case <-sigCtx.Done():
+				require.Equal(t, int32(0), atomic.LoadInt32(signalFired))
+			}
+		},
+	}
+
+	for name, f := range testcases {
+		t.Run(name, f)
+	}
+}
+
+func TestSleepCtx(t *testing.T) {
+	var (
+		veryShort = 1 * time.Microsecond
+		long      = 10 * time.Second
+	)
+
+	// make sure errors are reported correctly
+
+	ctx := context.Background()
+	err := sleepCtx(ctx, veryShort)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = sleepCtx(ctx, long)
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.Canceled)
+
+	ctx, cancel = context.WithTimeout(context.Background(), veryShort)
+	err = sleepCtx(ctx, long)
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	cancel()
+
+	ctx, cancel = context.WithTimeout(context.Background(), long)
+	cancel()
+	err = sleepCtx(ctx, long)
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.Canceled)
+}