Move version checks for rollbacks in update main function

Author: pchila

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

@@ -161,22 +161,17 @@ func markUpgradeProvider(updateActiveCommit updateActiveCommitFunc, writeFile wr
 			Details:           upgradeDetails,
 		}
 
-		if agent.parsedVersion != nil && !agent.parsedVersion.Less(*Version_9_2_0_SNAPSHOT) {
-			// if we have a not empty rollback window, write the prev version in the rollbacks_available field
-			// we also need to check the destination version because the manual rollback and delayed cleanup will be
-			// handled by that version of agent, so it needs to be recent enough
-			for _, rollback := range availableRollbacks {
-				rollbackAvailable := RollbackAvailable{
-					Version: rollback.Version,
-					Home:    rollback.VersionedHome,
-				}
-
-				if rollback.TTL != nil {
-					rollbackAvailable.ValidUntil = *rollback.TTL
-				}
-
-				marker.RollbacksAvailable = append(marker.RollbacksAvailable, rollbackAvailable)
+		for _, rollback := range availableRollbacks {
+			rollbackAvailable := RollbackAvailable{
+				Version: rollback.Version,
+				Home:    rollback.VersionedHome,
 			}
+
+			if rollback.TTL != nil {
+				rollbackAvailable.ValidUntil = *rollback.TTL
+			}
+
+			marker.RollbacksAvailable = append(marker.RollbacksAvailable, rollbackAvailable)
 		}
 
 		markerBytes, err := yaml.Marshal(newMarkerSerializer(marker))

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

@@ -183,61 +183,6 @@ func TestMarkUpgrade(t *testing.T) {
 				assert.Equal(t, expectedMarker, actualMarker)
 			},
 		},
-		{
-			name: "rollback window specified but new version is too low - no rollbacks",
-			args: args{
-				updatedOn: updatedOnNow,
-				currentAgent: agentInstall{
-					parsedVersion: parsed456SNAPSHOT,
-					version:       "4.5.6-SNAPSHOT",
-					hash:          "curagt",
-					versionedHome: filepath.Join("data", "elastic-agent-4.5.6-SNAPSHOT-curagt"),
-				},
-				previousAgent: agentInstall{
-					parsedVersion: parsed123SNAPSHOT,
-					version:       "1.2.3-SNAPSHOT",
-					hash:          "prvagt",
-					versionedHome: filepath.Join("data", "elastic-agent-1.2.3-SNAPSHOT-prvagt"),
-				},
-				action:  nil,
-				details: details.NewDetails("4.5.6-SNAPSHOT", details.StateReplacing, ""),
-				availableRollbacks: []v1.AgentInstallDesc{
-					{
-						OptionalTTLItem: v1.OptionalTTLItem{
-							TTL: &twentyFourHoursFromNow,
-						},
-						Version:       "1.2.3-SNAPSHOT",
-						Hash:          "prvagt",
-						VersionedHome: filepath.Join("data", "elastic-agent-1.2.3-SNAPSHOT-prvagt"),
-						Flavor:        "basic",
-						Active:        false,
-					},
-				},
-			},
-			wantErr: assert.NoError,
-			assertAfterMark: func(t *testing.T, dataDir string) {
-				actualMarker, err := LoadMarker(dataDir)
-				require.NoError(t, err, "error reading actualMarker content after writing")
-
-				expectedMarker := &UpdateMarker{
-					Version:           "4.5.6-SNAPSHOT",
-					Hash:              "curagt",
-					VersionedHome:     filepath.Join("data", "elastic-agent-4.5.6-SNAPSHOT-curagt"),
-					UpdatedOn:         updatedOnNow,
-					PrevVersion:       "1.2.3-SNAPSHOT",
-					PrevHash:          "prvagt",
-					PrevVersionedHome: filepath.Join("data", "elastic-agent-1.2.3-SNAPSHOT-prvagt"),
-					Acked:             false,
-					Action:            nil,
-					Details: &details.Details{
-						TargetVersion: "4.5.6-SNAPSHOT",
-						State:         "UPG_REPLACING",
-						ActionID:      "",
-					},
-				}
-				assert.Equal(t, expectedMarker, actualMarker)
-			},
-		},
 		{
 			name: "rollback window specified and new version is at least 9.2.0-SNAPSHOT - available rollbacks must be present",
 			args: args{

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

@@ -513,7 +513,7 @@ func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, s
 		versionedHome: currentVersionedHome,
 	}
 
-	availableRollbacks := getAvailableRollbacks(rollbackWindow, rotationTimestamp, unpackRes.VersionedHome, modifiedInstallDescriptor)
+	availableRollbacks := getAvailableRollbacks(rollbackWindow, parsedVersion, rotationTimestamp, unpackRes.VersionedHome, modifiedInstallDescriptor)
 
 	if err := u.markUpgrade(u.log,
 		paths.Data(), // data dir to place the marker in
@@ -554,12 +554,19 @@ func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, s
 	return cb, nil
 }
 
-func getAvailableRollbacks(rollbackWindow time.Duration, now time.Time, newVersionedHome string, descriptor *v1.InstallDescriptor) []v1.AgentInstallDesc {
+func getAvailableRollbacks(rollbackWindow time.Duration, version *agtversion.ParsedSemVer, now time.Time, newVersionedHome string, descriptor *v1.InstallDescriptor) []v1.AgentInstallDesc {
 	if rollbackWindow == 0 {
 		// if there's no rollback window it means that no rollback should survive the watcher cleanup at the end of the grace period.
 		return nil
 	}
 
+	if version == nil || version.Less(*Version_9_2_0_SNAPSHOT) {
+		// if we have a not empty rollback window, write the prev version in the rollbacks_available field
+		// we also need to check the destination version because the manual rollback and delayed cleanup will be
+		// handled by that version of agent, so it needs to be recent enough
+		return nil
+	}
+
 	res := make([]v1.AgentInstallDesc, 0, len(descriptor.AgentInstalls))
 	for _, installDesc := range descriptor.AgentInstalls {
 		if installDesc.VersionedHome != newVersionedHome && (installDesc.TTL == nil || now.Before(*installDesc.TTL)) {