Merge pull request #11560 from mardy/restore-conflict-nocleanup

overlord: execute snapshot cleanup in task

5 files changed, 39 insertions, 10 deletions

diff --git a/overlord/snapshotstate/snapshotmgr.go b/overlord/snapshotstate/snapshotmgr.go
index 0d596bfe82..f926629360 100644
--- a/overlord/snapshotstate/snapshotmgr.go
+++ b/overlord/snapshotstate/snapshotmgr.go
@@ -73,7 +73,7 @@ func Manager(st *state.State, runner *state.TaskRunner) *SnapshotManager {
	runner.AddHandler("forget-snapshot", doForget, nil)
	runner.AddHandler("check-snapshot", doCheck, nil)
	runner.AddHandler("restore-snapshot", doRestore, undoRestore)
-	runner.AddCleanup("restore-snapshot", cleanupRestore)
+	runner.AddHandler("cleanup-after-restore", doCleanupAfterRestore, nil)
 
	manager := &SnapshotManager{
	state: st,
@@ -371,6 +371,23 @@ func undoRestore(task *state.Task, _ *tomb.Tomb) error {
	return nil
 }
 
+func doCleanupAfterRestore(task *state.Task, tomb *tomb.Tomb) error {
+	st := task.State()
+	st.Lock()
+	restoreTasks := task.WaitTasks()
+	st.Unlock()
+	for _, t := range restoreTasks {
+	if err := cleanupRestore(t, tomb); err != nil {
+	logger.Noticef("Cleanup of restore task %s failed: %v", task.ID(), err)
+	// do not quit the loop: we must perform all cleanups anyway
+	}
+	}
+
+	// Also, do not return an error here: we don't want a failed cleanup to
+	// trigger an undo of the restore operation
+	return nil
+}
+
 func cleanupRestore(task *state.Task, _ *tomb.Tomb) error {
	var restoreState backend.RestoreState
 
diff --git a/overlord/snapshotstate/snapshotmgr_test.go b/overlord/snapshotstate/snapshotmgr_test.go
index f37d25711a..8b40d867f4 100644
--- a/overlord/snapshotstate/snapshotmgr_test.go
+++ b/overlord/snapshotstate/snapshotmgr_test.go
@@ -53,6 +53,7 @@ func (snapshotSuite) TestManager(c *check.C) {
	sort.Strings(kinds)
	c.Check(kinds, check.DeepEquals, []string{
	"check-snapshot",
+	"cleanup-after-restore",
	"forget-snapshot",
	"restore-snapshot",
	"save-snapshot",
diff --git a/overlord/snapshotstate/snapshotstate.go b/overlord/snapshotstate/snapshotstate.go
index 92eed2b07b..0cd153dbbf 100644
--- a/overlord/snapshotstate/snapshotstate.go
+++ b/overlord/snapshotstate/snapshotstate.go
@@ -522,6 +522,16 @@ func Restore(st *state.State, setID uint64, snapNames []string, users []string)
	ts.AddTask(task)
	}
 
+	if len(summaries) > 0 {
+	// take care of cleaning up all restore working state if all the
+	// restore tasks succeeded; if they didn't, the undo logic will take
+	// care of this
+	desc := fmt.Sprintf("Cleanup after restore from snapshot set #%d", setID)
+	task := st.NewTask("cleanup-after-restore", desc)
+	task.WaitAll(ts)
+	ts.AddTask(task)
+	}
+
	return snapsFound, ts, nil
 }
 
diff --git a/overlord/snapshotstate/snapshotstate_test.go b/overlord/snapshotstate/snapshotstate_test.go
index 3d62b21d13..8836c2f212 100644
--- a/overlord/snapshotstate/snapshotstate_test.go
+++ b/overlord/snapshotstate/snapshotstate_test.go
@@ -1118,8 +1118,9 @@ func (snapshotSuite) TestRestoreWorksWithCompatibleEpoch(c *check.C) {
	c.Assert(err, check.IsNil)
	c.Check(found, check.DeepEquals, []string{"a-snap"})
	tasks := taskset.Tasks()
-	c.Assert(tasks, check.HasLen, 1)
+	c.Assert(tasks, check.HasLen, 2)
	c.Check(tasks[0].Kind(), check.Equals, "restore-snapshot")
+	c.Check(tasks[1].Kind(), check.Equals, "cleanup-after-restore")
	c.Check(tasks[0].Summary(), check.Equals, `Restore data of snap "a-snap" from snapshot set #42`)
	var snapshot map[string]interface{}
	c.Check(tasks[0].Get("snapshot-setup", &snapshot), check.IsNil)
@@ -1154,8 +1155,9 @@ func (snapshotSuite) TestRestore(c *check.C) {
	c.Assert(err, check.IsNil)
	c.Check(found, check.DeepEquals, []string{"a-snap"})
	tasks := taskset.Tasks()
-	c.Assert(tasks, check.HasLen, 1)
+	c.Assert(tasks, check.HasLen, 2)
	c.Check(tasks[0].Kind(), check.Equals, "restore-snapshot")
+	c.Check(tasks[1].Kind(), check.Equals, "cleanup-after-restore")
	c.Check(tasks[0].Summary(), check.Equals, `Restore data of snap "a-snap" from snapshot set #42`)
	var snapshot map[string]interface{}
	c.Check(tasks[0].Get("snapshot-setup", &snapshot), check.IsNil)
@@ -1255,7 +1257,9 @@ func testRestoreIntegration(c *check.C, snapDataDir string, opts *dirs.SnapDirOp
	defer st.Unlock()
 
	// the three restores warn about the missing home (but no errors, no panics)
-	for _, task := range change.Tasks() {
+	c.Assert(change.Tasks(), check.HasLen, 4)
+	restoreTasks := change.Tasks()[:3]
+	for _, task := range restoreTasks {
	c.Check(strings.Join(task.Log(), "\n"), check.Matches, `.* Skipping restore of "[^"]+/home/b-user/[^"]+" as "[^"]+/home/b-user" doesn't exist.`)
	}
 
@@ -1333,8 +1337,9 @@ func (snapshotSuite) TestRestoreIntegrationFails(c *check.C) {
	defer st.Unlock()
 
	tasks := change.Tasks()
-	c.Check(tasks, check.HasLen, 3)
-	for _, task := range tasks {
+	c.Check(tasks, check.HasLen, 4)
+	restoreTasks := tasks[:3]
+	for _, task := range restoreTasks {
	if strings.Contains(task.Summary(), `"too-snap"`) {
	// too-snap was set up to fail, should always fail with
	// 'permission denied' (see the mkdirall w/mode 0 above)
diff --git a/tests/main/hidden-snap-dir/task.yaml b/tests/main/hidden-snap-dir/task.yaml
index fbd1ecbaae..575221430b 100644
--- a/tests/main/hidden-snap-dir/task.yaml
+++ b/tests/main/hidden-snap-dir/task.yaml
@@ -1,9 +1,5 @@
 summary: Check that the experimental hidden dir feature migrates the dir
 
-# this test is flaky on CentOS and openSUSE due to an issue w/ 'snap remove'
-# https://bugs.launchpad.net/snapd/+bug/1959036
-systems: [-centos-*, -opensuse-*]
-
 environment:
 NAME: test-snapd-tools