Support version in watcher cleanup rollback (#4176)

* extract directory paths to allow unit testing (instead of nested paths.* calls)
* introduce tests for watcher cleanup and rollback
* close data directory after cleanup

internal/pkg/agent/application/paths/common.go

@@ -180,11 +180,16 @@ func ExternalInputs() string {
 
 // Data returns the data directory for Agent
 func Data() string {
+	return DataFrom(Top())
+}
+
+// DataFrom returns the data directory for Agent using the passed directory as top path
+func DataFrom(topDirPath string) string {
 	if unversionedHome {
 		// unversioned means the topPath is the data path
-		return topPath
+		return topDirPath
 	}
-	return filepath.Join(Top(), "data")
+	return filepath.Join(topDirPath, "data")
 }
 
 // Run returns the run directory for Agent

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

@@ -33,64 +33,73 @@ const (
 )
 
 // Rollback rollbacks to previous version which was functioning before upgrade.
-func Rollback(ctx context.Context, log *logger.Logger, prevVersionedHome, prevHash string, currentHash string) error {
-	symlinkPath := filepath.Join(paths.Top(), agentName)
+func Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error {
+	symlinkPath := filepath.Join(topDirPath, agentName)
 
 	var symlinkTarget string
 	if prevVersionedHome != "" {
-		symlinkTarget = paths.BinaryPath(filepath.Join(paths.Top(), prevVersionedHome), agentName)
+		symlinkTarget = paths.BinaryPath(filepath.Join(topDirPath, prevVersionedHome), agentName)
 	} else {
 		// fallback for upgrades that didn't use the manifest and path remapping
 		hashedDir := fmt.Sprintf("%s-%s", agentName, prevHash)
 		// paths.BinaryPath properly derives the binary directory depending on the platform. The path to the binary for macOS is inside of the app bundle.
-		symlinkTarget = paths.BinaryPath(filepath.Join(paths.Top(), "data", hashedDir), agentName)
+		symlinkTarget = paths.BinaryPath(filepath.Join(paths.DataFrom(topDirPath), hashedDir), agentName)
 	}
 	// change symlink
-	if err := changeSymlinkInternal(log, symlinkPath, symlinkTarget); err != nil {
+	if err := changeSymlinkInternal(log, topDirPath, symlinkPath, symlinkTarget); err != nil {
 		return err
 	}
 
 	// revert active commit
-	if err := UpdateActiveCommit(log, prevHash); err != nil {
+	if err := UpdateActiveCommit(log, topDirPath, prevHash); err != nil {
 		return err
 	}
 
 	// Restart
 	log.Info("Restarting the agent after rollback")
-	if err := restartAgent(ctx, log); err != nil {
+	if err := restartAgent(ctx, log, c); err != nil {
 		return err
 	}
 
 	// cleanup everything except version we're rolling back into
-	return Cleanup(log, prevVersionedHome, prevHash, true, true)
+	return Cleanup(log, topDirPath, prevVersionedHome, prevHash, true, true)
 }
 
 // Cleanup removes all artifacts and files related to a specified version.
-func Cleanup(log *logger.Logger, currentVersionedHome, currentHash string, removeMarker bool, keepLogs bool) error {
+func Cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool) error {
 	log.Infow("Cleaning up upgrade", "hash", currentHash, "remove_marker", removeMarker)
 	<-time.After(afterRestartDelay)
 
+	// data directory path
+	dataDirPath := paths.DataFrom(topDirPath)
+
 	// remove upgrade marker
 	if removeMarker {
-		if err := CleanMarker(log); err != nil {
+		if err := CleanMarker(log, dataDirPath); err != nil {
 			return err
 		}
 	}
 
 	// remove data/elastic-agent-{hash}
-	dataDir, err := os.Open(paths.Data())
+	dataDir, err := os.Open(dataDirPath)
 	if err != nil {
 		return err
 	}
+	defer func(dataDir *os.File) {
+		err := dataDir.Close()
+		if err != nil {
+			log.Errorw("Error closing data directory", "file.directory", dataDirPath)
+		}
+	}(dataDir)
 
 	subdirs, err := dataDir.Readdirnames(0)
 	if err != nil {
 		return err
 	}
 
 	// remove symlink to avoid upgrade failures, ignore error
-	prevSymlink := prevSymlinkPath()
-	log.Infow("Removing previous symlink path", "file.path", prevSymlinkPath())
+	prevSymlink := prevSymlinkPath(topDirPath)
+	log.Infow("Removing previous symlink path", "file.path", prevSymlinkPath(topDirPath))
 	_ = os.Remove(prevSymlink)
 
 	dirPrefix := fmt.Sprintf("%s-", agentName)
@@ -113,7 +122,7 @@ func Cleanup(log *logger.Logger, currentVersionedHome, currentHash string, remov
 			continue
 		}
 
-		hashedDir := filepath.Join(paths.Data(), dir)
+		hashedDir := filepath.Join(dataDirPath, dir)
 		log.Infow("Removing hashed data directory", "file.path", hashedDir)
 		var ignoredDirs []string
 		if keepLogs {
@@ -155,9 +164,8 @@ func InvokeWatcher(log *logger.Logger) error {
 	return nil
 }
 
-func restartAgent(ctx context.Context, log *logger.Logger) error {
+func restartAgent(ctx context.Context, log *logger.Logger, c client.Client) error {
 	restartViaDaemonFn := func(ctx context.Context) error {
-		c := client.New()
 		connectCtx, connectCancel := context.WithTimeout(ctx, 3*time.Second)
 		defer connectCancel()
 		err := c.Connect(connectCtx, grpc.WithBlock(), grpc.WithDisableRetry())