Add a pre-restart hook to Rollback operation

Author: pchila

.mockery.yaml

@@ -30,13 +30,16 @@ packages:
     interfaces:
       Agent:
   github.com/elastic/elastic-agent/internal/pkg/agent/cmd:
+    config:
+      inpackage: True
+      with-expecter: True
+      dir: "{{.InterfaceDirRelative}}"
+      mockname: "{{.Mock}}{{.InterfaceName | firstUpper}}"
+      outpkg: "{{.PackageName}}"
+      filename: "{{.Mock | lower}}_{{.InterfaceName | lower}}_test.go"
     interfaces:
       agentWatcher:
-        config:
-          mockname: "AgentWatcher"
       installationModifier:
-        config:
-          mockname: "InstallationModifier"
   github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade:
     config:
       inpackage: True

internal/pkg/agent/application/upgrade/rollback.go

@@ -35,6 +35,36 @@ const (
 
 // Rollback rollbacks to previous version which was functioning before upgrade.
 func Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error {
+	return RollbackWithOpts(ctx, log, c, topDirPath, prevVersionedHome, prevHash)
+}
+
+var FatalRollbackError = errors.New("Fatal rollback error")
+
+type RollbackHook func(ctx context.Context, log *logger.Logger, topDirPath string) error
+type rollbackSettings struct {
+	preRestartHook RollbackHook
+}
+
+func newRollbackSettings(opts ...RollbackOpt) *rollbackSettings {
+	rs := new(rollbackSettings)
+	for _, opt := range opts {
+		opt(rs)
+	}
+	return rs
+}
+
+type RollbackOpt func(*rollbackSettings)
+
+func WithPreRestartHook(h RollbackHook) RollbackOpt {
+	return func(s *rollbackSettings) {
+		s.preRestartHook = h
+	}
+}
+
+func RollbackWithOpts(ctx context.Context, log *logger.Logger, c client.Client, topDirPath string, prevVersionedHome string, prevHash string, opts ...RollbackOpt) error {
+
+	settings := newRollbackSettings(opts...)
+
 	symlinkPath := filepath.Join(topDirPath, agentName)
 
 	var symlinkTarget string
@@ -56,6 +86,18 @@ func Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPa
 		return err
 	}
 
+	// Hook
+	if settings.preRestartHook != nil {
+		hookErr := settings.preRestartHook(ctx, log, topDirPath)
+		if hookErr != nil {
+			if errors.Is(hookErr, FatalRollbackError) {
+				return fmt.Errorf("pre-restart hook failed: %w", hookErr)
+			} else {
+				log.Warnf("pre-restart hook failed: %v", hookErr)
+			}
+		}
+	}
+
 	// Restart
 	log.Info("Restarting the agent after rollback")
 	if err := restartAgent(ctx, log, c); err != nil {

testing/mocks/internal_/pkg/agent/cmd/agent_watcher_mock.go renamed to internal/pkg/agent/cmd/mock_agentwatcher_test.go

@@ -86,9 +86,11 @@ type agentWatcher interface {
 	Watch(ctx context.Context, tilGrace, errorCheckInterval time.Duration, log *logp.Logger) error
 }
 
+type rollbackHook func(ctx context.Context, log *logger.Logger, topDirPath string) error
+
 type installationModifier interface {
 	Cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool) error
-	Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error
+	Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string, preRestart rollbackHook) error
 }
 
 func watchCmd(log *logp.Logger, topDir string, cfg *configuration.UpgradeWatcherConfig, watcher agentWatcher, installModifier installationModifier) error {
@@ -120,25 +122,29 @@ func watchCmd(log *logp.Logger, topDir string, cfg *configuration.UpgradeWatcher
 		_ = locker.Unlock()
 	}()
 
-	if marker.DesiredOutcome == upgrade.OUTCOME_ROLLBACK && marker.Details.State != details.StateRollback {
+	if marker.DesiredOutcome == upgrade.OUTCOME_ROLLBACK && marker.Details != nil && marker.Details.State != details.StateRollback {
 		// TODO: there should be some sanity check in rollback functions like the installation we are going back to should exist and work
 		log.Infof("rolling back because of DesiredOutcome=%s", marker.DesiredOutcome.String())
-		err = installModifier.Rollback(context.Background(), log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash)
-		if err != nil {
-			return fmt.Errorf("rolling back: %w", err)
-		}
 
-		if marker.Details == nil {
-			actionID := ""
-			if marker.Action != nil {
-				actionID = marker.Action.ActionID
+		updateMarkerAndDetails := func(_ context.Context, _ *logger.Logger, _ string) error {
+			if marker.Details == nil {
+				actionID := ""
+				if marker.Action != nil {
+					actionID = marker.Action.ActionID
+				}
+				marker.Details = details.NewDetails(marker.Version, details.StateRollback, actionID)
 			}
-			marker.Details = details.NewDetails(marker.Version, details.StateRollback, actionID)
+			marker.Details.SetStateWithReason(details.StateRollback, details.ReasonManualRollback)
+			err = upgrade.SaveMarker(dataDir, marker, true)
+			if err != nil {
+				return fmt.Errorf("saving marker after rolling back: %w", err)
+			}
+			return nil
 		}
-		marker.Details.SetStateWithReason(details.StateRollback, details.ReasonManualRollback)
-		err = upgrade.SaveMarker(dataDir, marker, true)
+
+		err = installModifier.Rollback(context.Background(), log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash, updateMarkerAndDetails)
 		if err != nil {
-			return fmt.Errorf("saving marker after rolling back: %w", err)
+			return fmt.Errorf("rolling back: %w", err)
 		}
 
 		return nil
@@ -180,7 +186,7 @@ func watchCmd(log *logp.Logger, topDir string, cfg *configuration.UpgradeWatcher
 		log.Error("Error detected, proceeding to rollback: %v", err)
 
 		upgradeDetails.SetStateWithReason(details.StateRollback, details.ReasonWatchFailed)
-		err = installModifier.Rollback(ctx, log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash)
+		err = installModifier.Rollback(ctx, log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash, nil)
 		if err != nil {
 			log.Error("rollback failed", err)
 			upgradeDetails.Fail(err)

internal/pkg/agent/cmd/mock_installationmodifier_test.go

@@ -29,8 +29,12 @@ func (a upgradeInstallationModifier) Cleanup(log *logger.Logger, topDirPath, cur
 	return upgrade.Cleanup(log, topDirPath, currentVersionedHome, currentHash, removeMarker, keepLogs)
 }
 
-func (a upgradeInstallationModifier) Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error {
-	return upgrade.Rollback(ctx, log, c, topDirPath, prevVersionedHome, prevHash)
+func (a upgradeInstallationModifier) Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string, preRestart rollbackHook) error {
+	var opts []upgrade.RollbackOpt
+	if preRestart != nil {
+		opts = append(opts, upgrade.WithPreRestartHook(upgrade.RollbackHook(preRestart)))
+	}
+	return upgrade.RollbackWithOpts(ctx, log, c, topDirPath, prevVersionedHome, prevHash, opts...)
 }
 
 func watch(ctx context.Context, tilGrace time.Duration, errorCheckInterval time.Duration, log *logger.Logger) error {