o/snapstate, o/devicestate: support components in snapstate.LinkNewBa…

…seOrKernel and snapstate.AddLinkNewBaseOrKernel

overlord/devicestate/devicestate_remodel_test.go

@@ -3867,14 +3867,17 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	for _, alreadyInstalledName := range []string{"pc-kernel-new", "core24-new", "pc-new"} {
 		snapYaml := "name: pc-kernel-new\nversion: 1\ntype: kernel\n"
 		channel := "other/edge"
+		rev := snap.R(222)
 		if alreadyInstalledName == "core24-new" {
 			snapYaml = "name: core24-new\nversion: 1\ntype: base\n"
+			rev = snap.R(223)
 		} else if alreadyInstalledName == "pc-new" {
 			snapYaml = "name: pc-new\nversion: 1\ntype: gadget\nbase: core24-new\n"
+			rev = snap.R(224)
 		}
 		si := &snap.SideInfo{
 			RealName: alreadyInstalledName,
-			Revision: snap.R(222),
+			Revision: rev,
 			SnapID:   snaptest.AssertedSnapID(alreadyInstalledName),
 		}
 		info := snaptest.MakeSnapFileAndDir(c, snapYaml, nil, si)

overlord/snapstate/snapstate.go

@@ -3404,10 +3404,41 @@ func LinkNewBaseOrKernel(st *state.State, name string, fromChange string) (*stat
 		InstanceKey: snapst.InstanceKey,
 	}
 
+	// note that prepare-snap doesn't actually do anything here, and is mostly
+	// used as a task to carry the snap-setup information.
 	prepareSnap := st.NewTask("prepare-snap", fmt.Sprintf(i18n.G("Prepare snap %q (%s) for remodel"), snapsup.InstanceName(), snapst.Current))
 	prepareSnap.Set("snap-setup", &snapsup)
-	prev := prepareSnap
+
 	ts := state.NewTaskSet(prepareSnap)
+	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
+	ts.MarkEdge(prepareSnap, SnapSetupEdge)
+
+	if err := addLinkNewBaseOrKernelTasks(st, snapst, ts, prepareSnap); err != nil {
+		return nil, err
+	}
+
+	return ts, nil
+}
+
+func addLinkNewBaseOrKernelTasks(st *state.State, snapst SnapState, ts *state.TaskSet, snapsupTask *state.Task) error {
+	tasks := ts.Tasks()
+	if len(tasks) == 0 {
+		return errors.New("internal error: task set must be seeded with at least one task")
+	}
+
+	prev := tasks[len(tasks)-1]
+	add := func(t *state.Task) {
+		t.Set("snap-setup-task", snapsupTask.ID())
+		t.WaitFor(prev)
+		ts.AddTask(t)
+		prev = t
+	}
+
+	info, err := snapst.CurrentInfo()
+	if err != nil {
+		return err
+	}
+
 	// preserve the same order as during the update
 	if info.Type() == snap.TypeKernel {
 		// this previously created a prepare-kernel-snap task, but this was not
@@ -3417,41 +3448,35 @@ func LinkNewBaseOrKernel(st *state.State, name string, fromChange string) (*stat
 		// installed, then it is expected that the drivers tree is present.
 
 		// kernel snaps can carry boot assets
-		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), snapsup.Type, snapsup.InstanceName(), snapst.Current))
-		gadgetUpdate.Set("snap-setup-task", prepareSnap.ID())
-		gadgetUpdate.WaitFor(prev)
-		ts.AddTask(gadgetUpdate)
-		prev = gadgetUpdate
+		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), info.Type(), info.InstanceName(), snapst.Current))
+		add(gadgetUpdate)
 	}
-	linkSnap := st.NewTask("link-snap", fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), snapsup.InstanceName(), snapst.Current))
-	linkSnap.Set("snap-setup-task", prepareSnap.ID())
-	linkSnap.WaitFor(prev)
-	ts.AddTask(linkSnap)
+
+	linkSnap := st.NewTask("link-snap", fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), info.InstanceName(), snapst.Current))
+	add(linkSnap)
 	ts.MarkEdge(linkSnap, MaybeRebootEdge)
 
-	for _, cs := range snapst.Sequence.ComponentsForRevision(snapst.Current) {
+	components := snapst.Sequence.ComponentsForRevision(snapst.Current)
+	compsupTasks := make([]string, 0, len(components))
+	for _, cs := range components {
 		compsup := ComponentSetup{
 			CompSideInfo: cs.SideInfo,
 			CompType:     cs.CompType,
 		}
 
-		compName := compsup.CompSideInfo.Component
+		cref := compsup.CompSideInfo.Component
 		compRev := compsup.CompSideInfo.Revision
 
-		prepare := st.NewTask("prepare-component", fmt.Sprintf(i18n.G("Prepare component %q (%s) for remodel"), compName, compRev))
-		prepare.Set("component-setup", compsup)
-		ts.AddTask(prepare)
+		link := st.NewTask("link-component", fmt.Sprintf(i18n.G("Make component %q (%s) available to the system during remodel"), cref, compRev))
+		link.Set("component-setup", compsup)
+		add(link)
 
-		link := st.NewTask("link-component", fmt.Sprintf(i18n.G("Make component %q (%s) available to the system during remodel"), compName, compRev))
-		link.Set("component-setup-task", prepare.ID())
-		link.WaitFor(prepare)
-		ts.AddTask(link)
+		compsupTasks = append(compsupTasks, link.ID())
 	}
 
-	// prepare-snap is the last task that carries no system modifications
-	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
-	ts.MarkEdge(prepareSnap, SnapSetupEdge)
-	return ts, nil
+	snapsupTask.Set("component-setup-tasks", compsupTasks)
+
+	return nil
 }
 
 func findSnapSetupTask(tasks []*state.Task) (*state.Task, *SnapSetup, error) {
@@ -3470,40 +3495,37 @@ func findSnapSetupTask(tasks []*state.Task) (*state.Task, *SnapSetup, error) {
 // AddLinkNewBaseOrKernel creates the same tasks as LinkNewBaseOrKernel but adds
 // them to the provided task set.
 func AddLinkNewBaseOrKernel(st *state.State, ts *state.TaskSet) (*state.TaskSet, error) {
+	if ts.MaybeEdge(LastBeforeLocalModificationsEdge) != nil {
+		return nil, errors.New("internal error: cannot add tasks to link new base or kernel to task set that introduces local modifications")
+	}
+
 	allTasks := ts.Tasks()
 	snapSetupTask, snapsup, err := findSnapSetupTask(allTasks)
 	if err != nil {
 		return nil, err
 	}
 	if snapSetupTask == nil {
-		return nil, fmt.Errorf("internal error: cannot identify task with snap-setup")
+		return nil, errors.New("internal error: cannot identify task with snap-setup")
 	}
-	// the first task added here waits for the last task in the existing set
-	prev := allTasks[len(allTasks)-1]
-	// preserve the same order as during the update
-	if snapsup.Type == snap.TypeKernel {
-		// kernel snaps can carry boot assets
-		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), snapsup.Type, snapsup.InstanceName(), snapsup.Revision()))
-		gadgetUpdate.Set("snap-setup-task", snapSetupTask.ID())
-		// wait for the last task in existing set
-		gadgetUpdate.WaitFor(prev)
-		ts.AddTask(gadgetUpdate)
-		prev = gadgetUpdate
-	}
-	linkSnap := st.NewTask("link-snap",
-		fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), snapsup.InstanceName(), snapsup.SideInfo.Revision))
-	linkSnap.Set("snap-setup-task", snapSetupTask.ID())
-	linkSnap.WaitFor(prev)
-	ts.AddTask(linkSnap)
-	ts.MarkEdge(linkSnap, MaybeRebootEdge)
-	// make sure that remodel can identify which tasks introduce actual
-	// changes to the system and order them correctly
-	if edgeTask := ts.MaybeEdge(LastBeforeLocalModificationsEdge); edgeTask == nil {
-		// no task in the task set is marked as last before system
-		// modifications are introduced, so we need to mark the last
-		// task in the set, as tasks introduced here modify system state
-		ts.MarkEdge(allTasks[len(allTasks)-1], LastBeforeLocalModificationsEdge)
+
+	var snapst SnapState
+	if err := Get(st, snapsup.InstanceName(), &snapst); err != nil {
+		return nil, err
 	}
+
+	if snapst.Current != snapsup.Revision() {
+		return nil, errors.New("internal error: cannot add tasks to link new base or kernel to task set that changes the snap revision")
+	}
+
+	// no task in the task set is marked as last before system modifications are
+	// introduced, so we need to mark the last task in the original set, as
+	// tasks introduced here modify system state
+	ts.MarkEdge(allTasks[len(allTasks)-1], LastBeforeLocalModificationsEdge)
+
+	if err := addLinkNewBaseOrKernelTasks(st, snapst, ts, snapSetupTask); err != nil {
+		return nil, err
+	}
+
 	return ts, nil
 }