o/snapstate: add some metadata to help with component remodeling

overlord/snapstate/component.go

@@ -92,14 +92,15 @@
 	// TODO:COMPS: verify validation sets here
 
 	snapsup := SnapSetup{
-		Base:        info.Base,
-		SideInfo:    &info.SideInfo,
-		Channel:     info.Channel,
-		Flags:       opts.Flags.ForSnapSetup(),
-		Type:        info.Type(),
-		Version:     info.Version,
-		PlugsOnly:   len(info.Slots) == 0,
-		InstanceKey: info.InstanceKey,
+		Base:                        info.Base,
+		SideInfo:                    &info.SideInfo,
+		Channel:                     info.Channel,
+		Flags:                       opts.Flags.ForSnapSetup(),
+		Type:                        info.Type(),
+		Version:                     info.Version,
+		PlugsOnly:                   len(info.Slots) == 0,
+		InstanceKey:                 info.InstanceKey,
+		ComponentExclusiveOperation: true,
 	}
 
 	setupSecurity := st.NewTask("setup-profiles",
@@ -143,6 +144,8 @@
 	setupSecurity.Set("component-setup-tasks", compSetupIDs)
 
 	ts := state.NewTaskSet(setupSecurity)
+	ts.MarkEdge(setupSecurity, SnapSetupEdge)
+
 	if kmodSetup != nil {
 		ts.AddTask(kmodSetup)
 	}
@@ -264,14 +267,15 @@
 	}
 
 	snapsup := SnapSetup{
-		Base:        info.Base,
-		SideInfo:    &info.SideInfo,
-		Channel:     info.Channel,
-		Flags:       opts.Flags.ForSnapSetup(),
-		Type:        info.Type(),
-		Version:     info.Version,
-		PlugsOnly:   len(info.Slots) == 0,
-		InstanceKey: info.InstanceKey,
+		Base:                        info.Base,
+		SideInfo:                    &info.SideInfo,
+		Channel:                     info.Channel,
+		Flags:                       opts.Flags.ForSnapSetup(),
+		Type:                        info.Type(),
+		Version:                     info.Version,
+		PlugsOnly:                   len(info.Slots) == 0,
+		InstanceKey:                 info.InstanceKey,
+		ComponentExclusiveOperation: true,
 	}
 	compSetup := ComponentSetup{
 		CompSideInfo: csi,
@@ -289,6 +293,19 @@
 	}
 
 	ts := componentTS.taskSet()
+
+	// TODO:COMPS: instead of doing this, we should convert this function to
+	// operate on multiple components so that it works like InstallComponents.
+	// this would improve performance, especially in the case of kernel module
+	// components.
+	begin, err := ts.Edge(BeginEdge)
+	if err != nil {
+		return nil, fmt.Errorf("internal error: cannot find begin edge on component install task set: %v", err)
+	}
+
+	begin.Set("component-setup-tasks", []string{componentTS.compSetupTaskID})
+	ts.MarkEdge(begin, SnapSetupEdge)
+
 	ts.JoinLane(generateLane(st, opts))
 
 	return ts, nil

overlord/snapstate/component_install_test.go

@@ -182,6 +182,20 @@ func verifyComponentInstallTasks(c *C, opts int, ts *state.TaskSet) {
 	} else {
 		c.Assert(t.Kind(), Equals, "prepare-component")
 	}
+
+	if opts&compOptMultiCompInstall == 0 {
+		snapsupTask, err := ts.Edge(snapstate.SnapSetupEdge)
+		c.Assert(err, IsNil)
+
+		var compsupsIDs []string
+		err = snapsupTask.Get("component-setup-tasks", &compsupsIDs)
+		c.Assert(err, IsNil)
+
+		// for now, all non-multi-component installs are by path, so this will
+		// point to prepare-component
+		c.Assert(snapsupTask.Kind(), Equals, "prepare-component")
+		c.Assert(compsupsIDs, DeepEquals, []string{snapsupTask.ID()})
+	}
 }
 
 func createTestComponent(c *C, snapName, compName string, snapInfo *snap.Info) (*snap.ComponentInfo, string) {
@@ -502,6 +516,7 @@ func (s *snapmgrTestSuite) TestInstallComponentPathForParallelInstall(c *C) {
 	var snapsup snapstate.SnapSetup
 	c.Assert(ts.Tasks()[0].Get("snap-setup", &snapsup), IsNil)
 	c.Assert(snapsup.InstanceKey, Equals, snapKey)
+	c.Assert(snapsup.ComponentExclusiveOperation, Equals, true)
 }
 
 func (s *snapmgrTestSuite) TestInstallComponentPathWrongSnap(c *C) {
@@ -1028,12 +1043,19 @@ func (s *snapmgrTestSuite) testInstallComponents(c *C, opts testInstallComponent
 	tss, err := snapstate.InstallComponents(context.Background(), s.state, components, info, nil, installOpts)
 	c.Assert(err, IsNil)
 
-	setupProfiles := tss[len(tss)-1].Tasks()[0]
+	setupTs := tss[len(tss)-1]
+
+	setupProfiles := setupTs.Tasks()[0]
 	c.Assert(setupProfiles.Kind(), Equals, "setup-profiles")
 
-	prepareKmodComps := tss[len(tss)-1].Tasks()[1]
+	prepareKmodComps := setupTs.Tasks()[1]
 	c.Assert(prepareKmodComps.Kind(), Equals, "prepare-kernel-modules-components")
 
+	snapsupTask, err := setupTs.Edge(snapstate.SnapSetupEdge)
+	c.Assert(err, IsNil)
+	c.Assert(snapsupTask.Kind(), Equals, "setup-profiles")
+	c.Assert(snapsupTask.Has("component-setup-tasks"), Equals, true)
+
 	expectedLane := opts.lane
 	if opts.transaction != "" && opts.lane == 0 {
 		expectedLane = 1
@@ -1052,6 +1074,7 @@ func (s *snapmgrTestSuite) testInstallComponents(c *C, opts testInstallComponent
 	snapsup, err := snapstate.TaskSnapSetup(prepareKmodComps)
 	c.Assert(err, IsNil)
 	c.Assert(snapsup, NotNil)
+	c.Assert(snapsup.ComponentExclusiveOperation, Equals, true)
 
 	for _, ts := range tss[0 : len(tss)-1] {
 		task := ts.Tasks()[0]

overlord/snapstate/snapmgr.go

@@ -159,6 +159,10 @@ type SnapSetup struct {
 	// case of an undo. Note that this cannot be tagged as omitempty, since we
 	// need to distinguish between empty and nil.
 	PreUpdateKernelModuleComponents []*snap.ComponentSideInfo `json:"pre-update-kernel-module-components"`
+
+	// ComponentExclusiveOperation is set if this SnapSetup exists only to deal with
+	// components, and not the snap itself.
+	ComponentExclusiveOperation bool `json:"component-exclusive-operation,omitempty"`
 }
 
 // ConfdbID identifies a confdb.

overlord/snapstate/snapstate.go

@@ -73,6 +73,7 @@
 
 const (
 	BeginEdge                        = state.TaskSetEdge("begin")
+	SnapSetupEdge                    = state.TaskSetEdge("snap-setup")
 	BeforeHooksEdge                  = state.TaskSetEdge("before-hooks")
 	HooksEdge                        = state.TaskSetEdge("hooks")
 	MaybeRebootEdge                  = state.TaskSetEdge("maybe-reboot")
@@ -786,6 +787,7 @@
 
 	installSet := state.NewTaskSet(tasks...)
 	installSet.MarkEdge(prereq, BeginEdge)
+	installSet.MarkEdge(prepare, SnapSetupEdge)
 	installSet.MarkEdge(setupAliases, BeforeHooksEdge)
 
 	// Let tasks know if they have to do something about restarts
@@ -1635,18 +1637,19 @@
 	}
 
 	snapsup := &SnapSetup{
-		Channel:            revOpts.Channel,
-		Base:               info.Base,
-		UserID:             opts.UserID,
-		Flags:              opts.Flags.ForSnapSetup(),
-		DownloadInfo:       &info.DownloadInfo,
-		SideInfo:           &info.SideInfo,
-		Type:               info.Type(),
-		Version:            info.Version,
-		InstanceKey:        info.InstanceKey,
-		CohortKey:          revOpts.CohortKey,
-		ExpectedProvenance: info.SnapProvenance,
-		DownloadBlobDir:    downloadDir,
+		Channel:                     revOpts.Channel,
+		Base:                        info.Base,
+		UserID:                      opts.UserID,
+		Flags:                       opts.Flags.ForSnapSetup(),
+		DownloadInfo:                &info.DownloadInfo,
+		SideInfo:                    &info.SideInfo,
+		Type:                        info.Type(),
+		Version:                     info.Version,
+		InstanceKey:                 info.InstanceKey,
+		CohortKey:                   revOpts.CohortKey,
+		ExpectedProvenance:          info.SnapProvenance,
+		DownloadBlobDir:             downloadDir,
+		ComponentExclusiveOperation: skipSnapDownload,
 	}
 
 	if sar.RedirectChannel != "" {
@@ -2434,22 +2437,30 @@
 		return nil, nil
 	}
 
-	var tasks []*state.Task
 	if err := checkChangeConflictIgnoringOneChange(st, snapst.InstanceName(), nil, opts.FromChange); err != nil {
 		return nil, err
 	}
 
+	var snapsupTask *state.Task
+
+	var tasks []*state.Task
 	if switchChannel || switchCohortKey {
 		summary := switchSummary(snapsup.InstanceName(), snapst.TrackingChannel, snapsup.Channel, snapst.CohortKey, snapsup.CohortKey)
 		switchSnap := st.NewTask("switch-snap-channel", summary)
 		switchSnap.Set("snap-setup", &snapsup)
+		snapsupTask = switchSnap
 
 		tasks = append(tasks, switchSnap)
 	}
 
 	if toggleIgnoreValidation {
 		toggle := st.NewTask("toggle-snap-flags", fmt.Sprintf(i18n.G("Toggle snap %q flags"), snapsup.InstanceName()))
-		toggle.Set("snap-setup", &snapsup)
+		if snapsupTask == nil {
+			toggle.Set("snap-setup", &snapsup)
+			snapsupTask = toggle
+		} else {
+			toggle.Set("snap-setup-task", snapsupTask.ID())
+		}
 
 		for _, tasks := range tasks {
 			toggle.WaitFor(tasks)
@@ -2458,7 +2469,12 @@
 		tasks = append(tasks, toggle)
 	}
 
-	return state.NewTaskSet(tasks...), nil
+	ts := state.NewTaskSet(tasks...)
+	if snapsupTask != nil {
+		ts.MarkEdge(snapsupTask, SnapSetupEdge)
+	}
+
+	return ts, nil
 }
 
 func splitEssentialUpdates(deviceCtx DeviceContext, updates []update) (essential, nonEssential []update) {
@@ -2805,7 +2821,10 @@
 	switchSnap := st.NewTask("switch-snap", summary)
 	switchSnap.Set("snap-setup", &snapsup)
 
-	return state.NewTaskSet(switchSnap), nil
+	ts := state.NewTaskSet(switchSnap)
+	ts.MarkEdge(switchSnap, SnapSetupEdge)
+
+	return ts, nil
 }
 
 // RevisionOptions control the selection of a snap revision.
@@ -3411,6 +3430,7 @@
 	ts.AddTask(linkSnap)
 	// prepare-snap is the last task that carries no system modifications
 	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
+	ts.MarkEdge(prepareSnap, SnapSetupEdge)
 	return ts, nil
 }
 
@@ -3534,6 +3554,7 @@
 	ts := state.NewTaskSet(prepareSnap, gadgetUpdate, gadgetCmdline)
 	// prepare-snap is the last task that carries no system modifications
 	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
+	ts.MarkEdge(prepareSnap, SnapSetupEdge)
 	return ts, nil
 }
 

overlord/snapstate/snapstate_install_test.go

@@ -237,6 +237,10 @@ func verifyInstallTasksWithComponents(c *C, typ snap.Type, opts, discards int, c
 		}
 	}
 
+	te, err := ts.Edge(snapstate.SnapSetupEdge)
+	c.Assert(err, IsNil)
+	c.Assert(te.Has("component-setup-tasks"), Equals, true)
+
 	var count int
 	for _, t := range ts.Tasks() {
 		switch t.Kind() {
@@ -251,6 +255,9 @@ func verifyInstallTasksWithComponents(c *C, typ snap.Type, opts, discards int, c
 			c.Assert(err, IsNil)
 			c.Check(compsup.CompSideInfo.Component.ComponentName, Equals, components[count])
 			count++
+		case "setup-profiles":
+			c.Assert(t.Has("component-setup-task"), Equals, false)
+			c.Assert(t.Has("component-setup"), Equals, false)
 		}
 	}
 	c.Check(count, Equals, len(components), Commentf("expected %d components, found %d", len(components), count))

overlord/snapstate/snapstate_test.go

@@ -872,6 +872,7 @@ func (s *snapmgrTestSuite) TestSwitchTasks(c *C) {
 
 	c.Assert(s.state.TaskCount(), Equals, len(ts.Tasks()))
 	c.Assert(taskKinds(ts.Tasks()), DeepEquals, []string{"switch-snap"})
+	c.Assert(ts.MaybeEdge(snapstate.SnapSetupEdge), NotNil)
 }
 
 func (s *snapmgrTestSuite) TestSwitchConflict(c *C) {
@@ -10176,7 +10177,8 @@ func (s *snapmgrTestSuite) TestDownloadWithComponents(c *C) {
 	c.Assert(begin, NotNil)
 	c.Check(begin.Kind(), Equals, "download-snap")
 
-	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir)
+	const componentExclusive = false
+	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir, componentExclusive)
 }
 
 func (s *snapmgrTestSuite) TestDownloadWithComponentsWithMismatchValidationSets(c *C) {
@@ -10388,7 +10390,8 @@ func (s *snapmgrTestSuite) TestDownloadWithComponentsWithValidationSets(c *C) {
 	c.Assert(begin, NotNil)
 	c.Check(begin.Kind(), Equals, "download-snap")
 
-	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir)
+	const componentExclusive = false
+	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir, componentExclusive)
 }
 
 func (s *snapmgrTestSuite) TestDownloadComponents(c *C) {
@@ -10463,10 +10466,11 @@ func (s *snapmgrTestSuite) TestDownloadComponents(c *C) {
 	c.Assert(begin, NotNil)
 	c.Check(begin.Kind(), Equals, "download-component")
 
-	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir)
+	const componentExclusive = true
+	verifySnapAndComponentSetupsForDownload(c, begin, ts, downloadDir, componentExclusive)
 }
 
-func verifySnapAndComponentSetupsForDownload(c *C, begin *state.Task, ts *state.TaskSet, downloadDir string) {
+func verifySnapAndComponentSetupsForDownload(c *C, begin *state.Task, ts *state.TaskSet, downloadDir string, componentExclusive bool) {
 	var snapsup snapstate.SnapSetup
 	err := begin.Get("snap-setup", &snapsup)
 	c.Assert(err, IsNil)
@@ -10482,6 +10486,8 @@ func verifySnapAndComponentSetupsForDownload(c *C, begin *state.Task, ts *state.
 		fmt.Sprintf("%s_%s.snap", snapsup.InstanceName(), snapsup.Revision()),
 	))
 
+	c.Assert(snapsup.ComponentExclusiveOperation, Equals, componentExclusive)
+
 	var compsupTaskIDs []string
 	err = begin.Get("component-setup-tasks", &compsupTaskIDs)
 	c.Assert(err, IsNil)

overlord/snapstate/snapstate_update_test.go

@@ -3117,6 +3117,7 @@ func (s *snapmgrTestSuite) TestUpdateSameRevisionSwitchesChannel(c *C) {
 	c.Assert(err, IsNil)
 	c.Check(ts.Tasks(), HasLen, 1)
 	c.Check(ts.Tasks()[0].Kind(), Equals, "switch-snap-channel")
+	c.Check(ts.MaybeEdge(snapstate.SnapSetupEdge).Kind(), Equals, "switch-snap-channel")
 }
 
 func (s *snapmgrTestSuite) TestUpdateSameRevisionSwitchesChannelConflict(c *C) {

overlord/snapstate/snapstatetest/snapstate.go

@@ -79,18 +79,21 @@ func InstallEssentialSnaps(c *check.C, st *state.State, base string, gadgetFiles
 		SnapID:   fakeSnapID("pc"),
 		Revision: snap.R(1),
 		RealName: "pc",
+		Channel:  "latest/stable",
 	}, InstallSnapOptions{Required: true})
 
 	InstallSnap(c, st, "name: pc-kernel\nversion: 1\ntype: kernel\n", nil, &snap.SideInfo{
 		SnapID:   fakeSnapID("pc-kernel"),
 		Revision: snap.R(1),
 		RealName: "pc-kernel",
+		Channel:  "latest/stable",
 	}, InstallSnapOptions{Required: true})
 
 	InstallSnap(c, st, fmt.Sprintf("name: %s\nversion: 1\ntype: base\n", base), nil, &snap.SideInfo{
 		SnapID:   fakeSnapID(base),
 		Revision: snap.R(1),
 		RealName: base,
+		Channel:  "latest/stable",
 	}, InstallSnapOptions{Required: true})
 
 	if bloader != nil {