From 015a2e4b7eef301a1fa4f4b6c4b945ad561946dc Mon Sep 17 00:00:00 2001
From: Andrew Phelps <136256549+andrewphelpsj@users.noreply.github.com>
Date: Tue, 2 Jul 2024 14:58:15 -0400
Subject: [PATCH] o/snapstate: account for remodeling when installing prereqs
 as well as updating them (#14137)

* o/snapstate: account for remodeling when installing prereqs as well as updating them

* o/snapstate: test that we do not wait for already installed prereqs that might be getting updated during a remodel

* o/snapstate: flatten code and update comment to account for refresh

* o/snapstate: move check for remodeling to be earlier, since we do not want to update prereqs during a remodel anyways
---
 overlord/snapstate/handlers.go             | 19 +++---
 overlord/snapstate/handlers_prereq_test.go | 70 ++++++++++++++++++++++
 2 files changed, 80 insertions(+), 9 deletions(-)

diff --git a/overlord/snapstate/handlers.go b/overlord/snapstate/handlers.go
index 02364781071..276640e6582 100644
--- a/overlord/snapstate/handlers.go
+++ b/overlord/snapstate/handlers.go
@@ -339,6 +339,15 @@ func (m *SnapManager) installOneBaseOrRequired(t *state.Task, snapName string, c
 		return true, nil
 	}
 
+	// if we are remodeling, then we should return early due to the way that
+	// tasks are ordered by the remodeling code. specifically, all snap
+	// downloads during a remodel happen prior to snap installation. thus,
+	// we cannot wait for snaps to be installed here. see remodelTasks for
+	// more information on how the tasks are ordered.
+	if deviceCtx.ForRemodeling() {
+		return nil, nil
+	}
+
 	if isInstalled {
 		if len(contentAttrs) > 0 {
 			// the default provider is already installed, update it if it's missing content attributes the snap needs
@@ -351,6 +360,7 @@ func (m *SnapManager) installOneBaseOrRequired(t *state.Task, snapName string, c
 		} else if ok {
 			return nil, onInFlight
 		}
+
 		return nil, nil
 	}
 
@@ -358,15 +368,6 @@ func (m *SnapManager) installOneBaseOrRequired(t *state.Task, snapName string, c
 	if ok, err := inProgress(snapName); err != nil {
 		return nil, err
 	} else if ok {
-		// if we are remodeling, then we should return early due to the way that
-		// tasks are ordered by the remodeling code. specifically, all snap
-		// downloads during a remodel happen prior to snap installation. thus,
-		// we cannot wait for snaps to be installed here. see remodelTasks for
-		// more information on how the tasks are ordered.
-		if deviceCtx.ForRemodeling() {
-			return nil, nil
-		}
-
 		return nil, onInFlight
 	}
 
diff --git a/overlord/snapstate/handlers_prereq_test.go b/overlord/snapstate/handlers_prereq_test.go
index b3e6f56b7f1..a6e52e1e951 100644
--- a/overlord/snapstate/handlers_prereq_test.go
+++ b/overlord/snapstate/handlers_prereq_test.go
@@ -403,6 +403,76 @@ func (s *prereqSuite) TestDoPrereqRetryWhenBaseInFlight(c *C) {
 	c.Check(chg.Status(), Equals, state.DoneStatus)
 }
 
+func (s *prereqSuite) TestDoPrereqNoRetryWhenBaseInFlightDuringRemodel(c *C) {
+	restore := snapstate.MockPrerequisitesRetryTimeout(1 * time.Millisecond)
+	defer restore()
+
+	s.runner.AddHandler("link-snap", func(task *state.Task, _ *tomb.Tomb) error {
+		st := task.State()
+		st.Lock()
+		defer st.Unlock()
+
+		snapsup, err := snapstate.TaskSnapSetup(task)
+		c.Assert(err, IsNil)
+		fmt.Println(snapsup.InstanceName())
+
+		return nil
+	}, nil)
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	restore = snapstatetest.MockDeviceContext(&snapstatetest.TrivialDeviceContext{
+		Remodeling: true,
+	})
+	defer restore()
+
+	// make it look like core is already installed
+	snapstate.Set(s.state, "core", &snapstate.SnapState{
+		Active: true,
+		Sequence: snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{
+			{RealName: "core", Revision: snap.R(1)},
+		}),
+		Current:  snap.R(1),
+		SnapType: "os",
+	})
+
+	// pretend foo gets installed and needs core (which we will make it look
+	// like the install is in progress)
+	prereqTask := s.state.NewTask("prerequisites", "foo")
+	prereqTask.Set("snap-setup", &snapstate.SnapSetup{
+		SideInfo: &snap.SideInfo{
+			RealName: "foo",
+		},
+	})
+
+	tCore := s.state.NewTask("link-snap", "Pretend core gets installed")
+	tCore.Set("snap-setup", &snapstate.SnapSetup{
+		SideInfo: &snap.SideInfo{
+			RealName: "core",
+			Revision: snap.R(11),
+		},
+	})
+
+	// this makes sure that the prereq task runs first, but the link-snap task
+	// for core will be found by the prereq task handler
+	tCore.WaitFor(prereqTask)
+
+	chg := s.state.NewChange("sample", "...")
+	chg.AddTask(prereqTask)
+	chg.AddTask(tCore)
+
+	s.state.Unlock()
+	s.se.Ensure()
+	s.se.Wait()
+	s.state.Lock()
+
+	// prereq task is done, but the core task isn't done. this means that we
+	// didn't wait for the core task to finish, since we are remodeling
+	c.Check(prereqTask.Status(), Equals, state.DoneStatus)
+	c.Check(tCore.Status(), Equals, state.DoStatus)
+}
+
 func (s *prereqSuite) TestDoPrereqChannelEnvvars(c *C) {
 	os.Setenv("SNAPD_BASES_CHANNEL", "edge")
 	defer os.Unsetenv("SNAPD_BASES_CHANNEL")