Fix daemon reload with backoff SIGINT handling (#4179)

* Fix daemon reload with backoff SIGINT handling

* Add unit tests

Author: aleksmaus

internal/pkg/agent/cmd/enroll_cmd.go

@@ -72,6 +72,9 @@ type enrollCmd struct {
 	remoteConfig remote.Config
 	agentProc    *process.Info
 	configPath   string
+
+	// For testability
+	daemonReloadFunc func(context.Context) error
 }
 
 // enrollCmdFleetServerOption define all the supported enrollment options for bootstrapping with Fleet Server.
@@ -192,10 +195,11 @@ func newEnrollCmdWithStore(
 	store saver,
 ) (*enrollCmd, error) {
 	return &enrollCmd{
-		log:         log,
-		options:     options,
-		configStore: store,
-		configPath:  configPath,
+		log:              log,
+		options:          options,
+		configStore:      store,
+		configPath:       configPath,
+		daemonReloadFunc: daemonReload,
 	}, nil
 }
 
@@ -462,17 +466,21 @@ func (c *enrollCmd) prepareFleetTLS() error {
 	return nil
 }
 
+const (
+	daemonReloadInitBackoff = time.Second
+	daemonReloadMaxBackoff  = time.Minute
+	daemonReloadRetries     = 5
+)
+
 func (c *enrollCmd) daemonReloadWithBackoff(ctx context.Context) error {
-	signal := make(chan struct{})
-	defer close(signal)
-	backExp := backoff.NewExpBackoff(signal, 1*time.Second, 1*time.Minute)
+	backExp := backoff.NewExpBackoff(ctx.Done(), daemonReloadInitBackoff, daemonReloadMaxBackoff)
 
 	var lastErr error
-	for i := 0; i < 5; i++ {
+	for i := 0; i < daemonReloadRetries; i++ {
 		attempt := i
 
 		c.log.Infof("Restarting agent daemon, attempt %d", attempt)
-		err := c.daemonReload(ctx)
+		err := c.daemonReloadFunc(ctx)
 		if err == nil {
 			return nil
 		}
@@ -486,14 +494,16 @@ func (c *enrollCmd) daemonReloadWithBackoff(ctx context.Context) error {
 		lastErr = err
 
 		c.log.Errorf("Restart attempt %d failed: '%s'. Waiting for %s", attempt, err, backExp.NextWait().String())
-		backExp.Wait()
-
+		// backoff Wait returns false if context.Done()
+		if !backExp.Wait() {
+			return ctx.Err()
+		}
 	}
 
 	return fmt.Errorf("could not reload agent's daemon, all retries failed. Last error: %w", lastErr)
 }
 
-func (c *enrollCmd) daemonReload(ctx context.Context) error {
+func daemonReload(ctx context.Context) error {
 	daemon := client.New()
 	err := daemon.Connect(ctx)
 	if err != nil {

internal/pkg/agent/cmd/enroll_cmd_test.go

@@ -463,6 +463,44 @@ func TestValidateEnrollFlags(t *testing.T) {
 	})
 }
 
+func TestDaemonReloadWithBackoff(t *testing.T) {
+	log, _ := logger.New("tst", false)
+
+	ctx, cn := context.WithCancel(context.Background())
+	// Cancel context
+	cn()
+
+	tests := []struct {
+		name             string
+		daemonReloadFunc func(ctx context.Context) error
+		wantErr          error
+	}{
+		{
+			name:             "daemonReloadSucceeded",
+			daemonReloadFunc: func(ctx context.Context) error { return nil },
+		},
+		{
+			name: "retryWithContextCancelled",
+			daemonReloadFunc: func(ctx context.Context) error {
+				return errors.New("failed") // Return some (not context's) error so it retries
+			},
+			wantErr: context.Canceled,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			cmd := enrollCmd{
+				log:              log,
+				daemonReloadFunc: tc.daemonReloadFunc,
+			}
+
+			err := cmd.daemonReloadWithBackoff(ctx)
+			assert.ErrorIs(t, err, tc.wantErr)
+		})
+	}
+}
+
 func withServer(
 	m func(t *testing.T) *http.ServeMux,
 	test func(t *testing.T, host string),