Fix managed rollback test (#11218)

* Add logging to stop/start operation for restartAgentNTimes in Test*UpgradeRollbackOnRestarts

* Do not reset lastPid at every iteration

* Switch to restarting based on version in TestFleetManagedUpgradeRollbackOnRestarts

Move from a fixed number of restarts to restarting only if the agent
matches the target version in the managed upgrade rollback integration
test.
This is done to avoid +-1 errors restarting the agent to simulate only
the upgrade agent crashing, while we can keep restarting for a given
time (controlled by a context with timeout) if the agent matches a
particular version. This should make the test more reliable and avoid
weird interaction between the restart timing and the watcher polling.

Author: pchila

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

@@ -86,6 +86,7 @@ func (ch *AgentWatcher) Run(ctx context.Context) {
 
 	ch.connectCounter = 0
 	ch.lostCounter = 0
+	ch.lastPid = -1
 
 	// tracking of an error runs in a separate goroutine, because
 	// the call to `watch.Recv` blocks and a timer is needed
@@ -143,7 +144,6 @@ func (ch *AgentWatcher) Run(ctx context.Context) {
 
 LOOP:
 	for {
-		ch.lastPid = -1
 		connectTimer := time.NewTimer(ch.checkInterval)
 		select {
 		case <-ctx.Done():

pkg/testing/fixture.go

@@ -55,6 +55,7 @@ type Fixture struct {
 	srcPackage string
 	workDir    string
 	extractDir string
+	socketPath string
 
 	installed   bool
 	installOpts *InstallOpts
@@ -186,12 +187,21 @@ func NewFixture(t *testing.T, version string, opts ...FixtureOpt) (*Fixture, err
 }
 
 // Client returns the Elastic Agent communication client.
+// This client is shared across multiple calls to Client()
 func (f *Fixture) Client() client.Client {
 	f.cMx.RLock()
 	defer f.cMx.RUnlock()
 	return f.c
 }
 
+// NewClient returns a new Elastic Agent communication client.
+// The client is NOT shared across multiple calls but a new instance is allocated every time
+func (f *Fixture) NewClient() client.Client {
+	f.cMx.Lock()
+	defer f.cMx.Unlock()
+	return client.New(client.WithAddress(f.socketPath))
+}
+
 // Version returns the Elastic Agent version.
 func (f *Fixture) Version() string {
 	return f.version
@@ -1136,6 +1146,12 @@ func (f *Fixture) prepareComponents(workDir string, components ...UsableComponen
 	return nil
 }
 
+func (f *Fixture) setSocketPath(socketPath string) {
+	f.cMx.Lock()
+	defer f.cMx.Unlock()
+	f.socketPath = socketPath
+}
+
 func (f *Fixture) setClient(c client.Client) {
 	f.cMx.Lock()
 	defer f.cMx.Unlock()
@@ -1605,6 +1621,7 @@ type AgentBinaryVersion struct {
 	Commit    string `yaml:"commit"`
 	BuildTime string `yaml:"build_time"`
 	Snapshot  bool   `yaml:"snapshot"`
+	Fips      bool   `yaml:"fips"`
 }
 
 // String returns the version string.

pkg/testing/fixture_install.go

@@ -278,6 +278,7 @@ func (f *Fixture) installNoPkgManager(ctx context.Context, installOpts *InstallO
 		socketPath = paths.ControlSocketFromPath(runtime.GOOS, f.workDir)
 	}
 	c := client.New(client.WithAddress(socketPath))
+	f.setSocketPath(socketPath)
 	f.setClient(c)
 
 	f.t.Cleanup(func() {

testing/integration/ess/upgrade_rollback_test.go

@@ -529,9 +529,13 @@ func managedRollbackRestartTest(ctx context.Context, t *testing.T, info *define.
 	// we expect ErrSkipGrace at this point, meaning that we finished installing but didn't wait for agent to become healthy
 	require.ErrorIs(t, err, ErrSkipGrace, "managed upgrade failed with unexpected error")
 
-	// A few seconds after the upgrade, deliberately restart upgraded Agent a
-	// couple of times to simulate Agent crashing.
-	restartAgentNTimes(t, 3, 10*time.Second)
+	installedAgentClient := from.NewClient()
+	targetVersion, err := to.ExecVersion(ctx)
+	require.NoError(t, err, "failed to get target version")
+	restartContext, cancel := context.WithTimeout(t.Context(), 1*time.Minute)
+	defer cancel()
+	// restart the agent only if it matches the (upgraded) target version
+	restartAgentVersion(restartContext, t, installedAgentClient, targetVersion.Binary, 10*time.Second)
 
 	// wait for the agent to be healthy and correct version
 	err = upgradetest.WaitHealthyAndVersion(ctx, from, startVersionInfo.Binary, 2*time.Minute, 10*time.Second, t)
@@ -657,42 +661,95 @@ func restartAgentNTimes(t *testing.T, noOfRestarts int, sleepBetweenIterations t
 
 	for restartIdx := 0; restartIdx < noOfRestarts; restartIdx++ {
 		time.Sleep(sleepBetweenIterations)
+		restartAgent(t, topPath, 5*time.Minute)
+	}
+}
 
-		t.Logf("Stopping agent via service to simulate crashing")
-		err := install.StopService(topPath, install.DefaultStopTimeout, install.DefaultStopInterval)
-		if err != nil && runtime.GOOS == define.Windows && strings.Contains(err.Error(), "The service has not been started.") {
-			// Due to the quick restarts every 10 seconds its possible that this is faster than Windows
-			// can handle. Decrementing restartIdx means that the loop will occur again.
-			t.Logf("Got an allowed error on Windows: %s", err)
-			err = nil
+func restartAgent(t *testing.T, topPath string, operationTimeout time.Duration) {
+	t.Logf("Stopping agent via service to simulate crashing")
+	stopRequested := time.Now()
+	err := install.StopService(topPath, install.DefaultStopTimeout, install.DefaultStopInterval)
+	if err != nil && runtime.GOOS == define.Windows && strings.Contains(err.Error(), "The service has not been started.") {
+		// Due to the quick restarts every sleepBetweenIterations its possible that this is faster than Windows
+		// can handle. Decrementing restartIdx means that the loop will occur again.
+		t.Logf("Got an allowed error on Windows: %s", err)
+		err = nil
+	}
+	require.NoError(t, err)
+
+	// ensure that it's stopped before starting it again
+	var status service.Status
+	var statusErr error
+	require.Eventuallyf(t, func() bool {
+		status, statusErr = install.StatusService(topPath)
+		if statusErr != nil {
+			return false
 		}
-		require.NoError(t, err)
+		return status != service.StatusRunning
+	}, operationTimeout, 500*time.Millisecond, "service never fully stopped (status: %v): %s", status, statusErr)
+	t.Logf("Stopped agent via service. Took roughly %s", time.Since(stopRequested))
+
+	// start it again
+	t.Logf("Starting agent via service to simulate crashing")
+	startRequested := time.Now()
+	err = install.StartService(topPath)
+	require.NoError(t, err)
 
-		// ensure that it's stopped before starting it again
-		var status service.Status
-		var statusErr error
-		require.Eventuallyf(t, func() bool {
-			status, statusErr = install.StatusService(topPath)
-			if statusErr != nil {
-				return false
-			}
-			return status != service.StatusRunning
-		}, 5*time.Minute, 1*time.Second, "service never fully stopped (status: %v): %s", status, statusErr)
-		t.Logf("Stopped agent via service to simulate crashing")
+	// ensure that it's started before next loop
+	require.Eventuallyf(t, func() bool {
+		status, statusErr = install.StatusService(topPath)
+		if statusErr != nil {
+			return false
+		}
+		return status == service.StatusRunning
+	}, operationTimeout, 500*time.Millisecond, "service never fully started (status: %v): %s", status, statusErr)
+	t.Logf("Started agent after stopping to simulate crashing. Took roughly %s", time.Since(startRequested))
+}
 
-		// start it again
-		t.Logf("Starting agent via service to simulate crashing")
-		err = install.StartService(topPath)
-		require.NoError(t, err)
+func restartAgentVersion(ctx context.Context, t *testing.T, client client.Client, targetVersion atesting.AgentBinaryVersion, restartInterval time.Duration) {
+	topPath := paths.Top()
+
+	ticker := time.NewTicker(restartInterval)
+	defer ticker.Stop()
 
-		// ensure that it's started before next loop
-		require.Eventuallyf(t, func() bool {
-			status, statusErr = install.StatusService(topPath)
-			if statusErr != nil {
-				return false
+	for {
+		select {
+		case <-ctx.Done():
+			t.Log("restart context is done, returning")
+			return
+
+		case <-ticker.C:
+			if !versionMatch(ctx, t, client, targetVersion) {
+				// version of running agent does not match the target, continue to the next iteration
+				continue
 			}
-			return status == service.StatusRunning
-		}, 5*time.Minute, 1*time.Second, "service never fully started (status: %v): %s", status, statusErr)
-		t.Logf("Started agent via service to simulate crashing")
+
+			restartAgent(t, topPath, restartInterval)
+		}
+
+	}
+}
+
+func versionMatch(ctx context.Context, t *testing.T, c client.Client, targetVersion atesting.AgentBinaryVersion) bool {
+	err := c.Connect(ctx)
+	if err != nil {
+		t.Logf("failed to connect to agent: %v", err)
+		return false
+	}
+	defer c.Disconnect()
+
+	actualVersion, err := c.Version(ctx)
+	if err != nil {
+		t.Logf("failed to detect agent version: %v", err)
+		return false
+	}
+
+	if actualVersion.Version != targetVersion.Version ||
+		actualVersion.Snapshot != targetVersion.Snapshot ||
+		actualVersion.Commit != targetVersion.Commit ||
+		actualVersion.Fips != targetVersion.Fips {
+		t.Logf("actual agent version %+v does not match target agent version %+v, skipping restart", actualVersion, targetVersion)
+		return false
 	}
+	return true
 }