o/snapstate: do not restart again for snapd along the undo path insid…

…e undoUnlinkCurrentSnap (canonical#14917)

* o/snapstate: do not reboot again for snapd along the undo path inside undoUnlinkCurrentSnap, actually we skip the entire link code for snapd  as this has already been done in undoLinkSnap. Add a wait for restart in undoRemoveAlias

* o/snapstate: fix the comment being wrong

* o/snapstate: fix wording of one of the comments that did not have clear intent

overlord/snapstate/handlers.go

@@ -1606,6 +1606,17 @@ func (m *SnapManager) undoUnlinkCurrentSnap(t *state.Task, _ *tomb.Tomb) error {
 
 	snapst.Active = true
 
+	// For snapd, we've already relinked the previous snapd (in undoLinkSnap)
+	// and restarted into that version of snapd at this point, so avoid redoing
+	// that which would have no effect.
+	if oldInfo.Type() == snap.TypeSnapd {
+		// mark as active again
+		Set(st, snapsup.InstanceName(), snapst)
+		return nil
+	}
+
+	// For all other snaps, including snapd bundled with the core snap,
+	// we must undo the unlinking of the old revision.
 	opts, err := SnapServiceOptions(st, oldInfo, nil)
 	if err != nil {
 		return err
@@ -1634,9 +1645,6 @@ func (m *SnapManager) undoUnlinkCurrentSnap(t *state.Task, _ *tomb.Tomb) error {
 
 	// if we just put back a previous a core snap, request a restart
 	// so that we switch executing its snapd
-	// TODO: This needs changing, right now we restart snapd after undoing 'link-snap',
-	// inside 'undoSetupProfiles() and even though we request a restart here we are
-	// not actually waiting for the restart to occur in the task that comes before this.
 	return m.finishTaskWithMaybeRestart(t, state.UndoneStatus, restartPossibility{info: oldInfo, RebootInfo: reboot})
 }
 
@@ -3985,6 +3993,13 @@ func (m *SnapManager) undoRemoveAliases(t *state.Task, _ *tomb.Tomb) error {
 		return err
 	}
 
+	// The previous task's undo (unlink-current-snap) may have triggered a restart
+	// so if that is the case ensure we wait for it to happen here.
+	logger.Debugf("finish restart from undoRemoveAliases")
+	if err := FinishRestart(t, snapsup, FinishRestartOptions{}); err != nil {
+		return err
+	}
+
 	if !snapst.AliasesPending {
 		// do nothing
 		return nil

overlord/snapstate/handlers_link_test.go

@@ -28,6 +28,7 @@ import (
 	"time"
 
 	. "gopkg.in/check.v1"
+	"gopkg.in/tomb.v2"
 
 	"github.com/snapcore/snapd/asserts"
 	"github.com/snapcore/snapd/boot"
@@ -772,6 +773,85 @@ func (s *linkSnapSuite) TestDoUnlinkCurrentSnapSnapdNop(c *C) {
 		}})
 }
 
+func (s *linkSnapSuite) TestDoUnlinkCurrentSnapNoRestartSnapd(c *C) {
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	s.runner.AddHandler("failure-task", func(_ *state.Task, _ *tomb.Tomb) error {
+		return fmt.Errorf("oh no!")
+	}, nil)
+
+	si := &snap.SideInfo{
+		RealName: "snapd",
+		Revision: snap.R(20),
+	}
+	snapstate.Set(s.state, "snapd", &snapstate.SnapState{
+		Sequence: snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{si}),
+		Current:  si.Revision,
+		Active:   true,
+	})
+
+	chg := s.state.NewChange("sample", "...")
+
+	raTask := s.state.NewTask("remove-aliases", "")
+	raTask.Set("snap-setup", &snapstate.SnapSetup{
+		SideInfo: si,
+		Channel:  "beta",
+	})
+	chg.AddTask(raTask)
+
+	unlinkTask := s.state.NewTask("unlink-current-snap", "")
+	unlinkTask.Set("snap-setup", &snapstate.SnapSetup{
+		SideInfo: si,
+		Channel:  "beta",
+	})
+	unlinkTask.WaitFor(raTask)
+	chg.AddTask(unlinkTask)
+
+	rmpTask := s.state.NewTask("failure-task", "")
+	rmpTask.WaitFor(unlinkTask)
+	chg.AddTask(rmpTask)
+
+	// Run the tasks we created
+	s.state.Unlock()
+	for i := 0; i < 5; i++ {
+		s.se.Ensure()
+		s.se.Wait()
+	}
+	s.state.Lock()
+
+	// And observe the results.
+	var snapst snapstate.SnapState
+	err := snapstate.Get(s.state, "snapd", &snapst)
+	c.Assert(err, IsNil)
+	c.Check(snapst.Active, Equals, true)
+	c.Check(snapst.Sequence.Revisions, HasLen, 1)
+	c.Check(snapst.Current, Equals, snap.R(20))
+	c.Check(raTask.Status(), Equals, state.UndoneStatus)
+	c.Check(unlinkTask.Status(), Equals, state.UndoneStatus)
+	c.Check(rmpTask.Status(), Equals, state.ErrorStatus)
+	// backend was called to unlink the snap
+	expected := fakeOps{
+		{
+			op:   "remove-snap-aliases",
+			name: "snapd",
+		},
+		{
+			op:          "run-inhibit-snap-for-unlink",
+			name:        "snapd",
+			inhibitHint: "refresh",
+		},
+		{
+			op: "update-aliases",
+		},
+	}
+	c.Check(s.fakeBackend.ops, DeepEquals, expected)
+
+	// no restarts must have been requested by 'unlink-current-snap'
+	c.Check(s.restartRequested, HasLen, 0)
+	c.Assert(unlinkTask.Log(), HasLen, 0)
+}
+
 func (s *linkSnapSuite) TestDoUnlinkSnapdUnlinks(c *C) {
 	s.state.Lock()
 	defer s.state.Unlock()
@@ -2449,7 +2529,6 @@ func (s *linkSnapSuite) TestUndoLinkSnapdNthInstall(c *C) {
 
 	// 1 restart from link snap
 	// 1 restart from undo of 'setup-profiles'
-	// 1 in undoUnlinkCurrentSnap (not tested here)
 	c.Check(s.restartRequested, DeepEquals, []restart.RestartType{restart.RestartDaemon, restart.RestartDaemon})
 	c.Assert(t.Log(), HasLen, 2)
 	c.Check(t.Log()[0], Matches, `.*INFO Requested daemon restart \(snapd snap\)\.`)