o/snapstate: enable installing additional components during a refresh

overlord/snapstate/snapstate.go

@@ -125,7 +125,7 @@ func (ins installSnapInfo) SnapBase() string {
 }
 
 func (ins installSnapInfo) Prereq(st *state.State, prqt PrereqTracker) []string {
-	return getKeys(defaultProviderContentAttrs(st, ins.Info, prqt))
+	return keys(defaultProviderContentAttrs(st, ins.Info, prqt))
 }
 
 // InsufficientSpaceError represents an error where there is not enough disk
@@ -1243,16 +1243,6 @@ func defaultProviderContentAttrs(st *state.State, info *snap.Info, prqt PrereqTr
 	return prqt.MissingProviderContentTags(info, repo)
 }
 
-func getKeys(kv map[string][]string) []string {
-	keys := make([]string, 0, len(kv))
-
-	for key := range kv {
-		keys = append(keys, key)
-	}
-
-	return keys
-}
-
 // validateFeatureFlags validates the given snap only uses experimental
 // features that are enabled by the user.
 func validateFeatureFlags(st *state.State, info *snap.Info) error {
@@ -1811,6 +1801,22 @@ func ResolveValidationSetsEnforcementError(ctx context.Context, st *state.State,
 	return tasksets, affected, nil
 }
 
+func keys[K comparable, V any](m map[K]V) []K {
+	keys := make([]K, 0, len(m))
+	for k := range m {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
+func unique[T comparable](s []T) []T {
+	m := make(map[T]struct{}, len(s))
+	for _, v := range s {
+		m[v] = struct{}{}
+	}
+	return keys(m)
+}
+
 // updateFilter is the type of function that can be passed to
 // updateManyFromChange so it filters the updates.
 //

overlord/snapstate/snapstate_update_test.go

@@ -15143,10 +15143,32 @@ func (s *snapmgrTestSuite) TestUpdateWithComponentsRunThroughLoseComponentsUndo(
 	})
 }
 
+func (s *snapmgrTestSuite) TestUpdateWithComponentsRunThroughAdditionalComponents(c *C) {
+	s.testUpdateWithComponentsRunThrough(c, updateWithComponentsOpts{
+		instanceKey:           "key",
+		components:            []string{"standard-component"},
+		postRefreshComponents: []string{"standard-component", "kernel-modules-component"},
+		additionalComponents:  []string{"kernel-modules-component"},
+		refreshAppAwarenessUX: true,
+	})
+}
+
+func (s *snapmgrTestSuite) TestUpdateWithComponentsRunThroughAdditionalComponentsUndo(c *C) {
+	s.testUpdateWithComponentsRunThrough(c, updateWithComponentsOpts{
+		instanceKey:           "key",
+		components:            []string{"standard-component"},
+		postRefreshComponents: []string{"standard-component", "kernel-modules-component"},
+		additionalComponents:  []string{"kernel-modules-component"},
+		refreshAppAwarenessUX: true,
+		undo:                  true,
+	})
+}
+
 type updateWithComponentsOpts struct {
 	instanceKey           string
 	components            []string
 	postRefreshComponents []string
+	additionalComponents  []string
 	refreshAppAwarenessUX bool
 	undo                  bool
 	useSameSnapRev        bool
@@ -15301,13 +15323,21 @@ func (s *snapmgrTestSuite) testUpdateWithComponentsRunThrough(c *C, opts updateW
 		InstanceKey:     opts.instanceKey,
 	})
 
-	var revOpts *snapstate.RevisionOptions
+	var revOpts snapstate.RevisionOptions
 	if opts.useSameSnapRev {
-		revOpts = &snapstate.RevisionOptions{Revision: currentSnapRev}
+		revOpts.Revision = currentSnapRev
 	}
 
-	ts, err := snapstate.Update(s.state, instanceName, revOpts, s.user.ID, snapstate.Flags{
-		NoReRefresh: true,
+	ts, err := snapstate.UpdateOne(context.Background(), s.state, snapstate.StoreUpdateGoal(snapstate.StoreUpdate{
+		InstanceName:         instanceName,
+		RevOpts:              revOpts,
+		AdditionalComponents: opts.additionalComponents,
+	}), nil, snapstate.Options{
+		Flags: snapstate.Flags{
+			NoReRefresh: true,
+			Transaction: client.TransactionPerSnap,
+		},
+		UserID: s.user.ID,
 	})
 	c.Assert(err, IsNil)
 

overlord/snapstate/storehelpers.go

@@ -522,7 +522,7 @@ func storeUpdatePlanCore(
 			return updatePlan{}, snap.NotInstalledError{Snap: update.InstanceName}
 		}
 
-		if snapst.HasActiveComponents() {
+		if snapst.HasActiveComponents() || len(update.AdditionalComponents) > 0 {
 			requestComponentsFromStore = true
 		}
 	}
@@ -558,7 +558,7 @@ func storeUpdatePlanCore(
 	}
 
 	for _, name := range localAmends {
-		hasLocalRevision[allSnaps[name]] = updates[name].RevOpts
+		hasLocalRevision[name] = allSnaps[name]
 	}
 
 	for id, actions := range amendActionsByUserID {
@@ -572,7 +572,7 @@ func storeUpdatePlanCore(
 	}
 
 	for _, name := range noStoreUpdates {
-		hasLocalRevision[allSnaps[name]] = updates[name].RevOpts
+		hasLocalRevision[name] = allSnaps[name]
 	}
 
 	for _, sar := range sars {
@@ -598,6 +598,10 @@ func storeUpdatePlanCore(
 			compNames = append(compNames, comp.Component.ComponentName)
 		}
 
+		// add the additional components that the caller requested to be
+		// installed
+		compNames = unique(append(compNames, up.AdditionalComponents...))
+
 		// compTargets will be filtered down to only the components that appear
 		// in the action result, meaning that we might install fewer components
 		// than we have installed right now
@@ -625,10 +629,15 @@ func storeUpdatePlanCore(
 
 	// consider snaps that already have a local copy of the revision that we are
 	// trying to install, skipping a trip to the store
-	for snapst, revOpts := range hasLocalRevision {
+	for name, snapst := range hasLocalRevision {
+		up, ok := updates[name]
+		if !ok {
+			return updatePlan{}, fmt.Errorf("internal error: unexpected update to local revision: %q", snapst.InstanceName())
+		}
+
 		var si *snap.SideInfo
-		if !revOpts.Revision.Unset() {
-			si = snapst.Sequence.Revisions[snapst.LastIndex(revOpts.Revision)].Snap
+		if !up.RevOpts.Revision.Unset() {
+			si = snapst.Sequence.Revisions[snapst.LastIndex(up.RevOpts.Revision)].Snap
 		} else {
 			si = snapst.CurrentSideInfo()
 		}
@@ -648,33 +657,37 @@ func storeUpdatePlanCore(
 			return updatePlan{}, err
 		}
 
-		compsups, err := componentSetupsForInstall(ctx, st, compsToInstall, *snapst, si.Revision, revOpts.Channel, opts)
+		// add the additional components that the caller requested to be
+		// installed
+		compsToInstall = unique(append(compsToInstall, up.AdditionalComponents...))
+
+		compsups, err := componentSetupsForInstall(ctx, st, compsToInstall, *snapst, si.Revision, up.RevOpts.Channel, opts)
 		if err != nil {
 			return updatePlan{}, err
 		}
 
 		// this must happen after the call to componentSetupsForInstall, since
 		// we can't set the channel to the tracking channel if we don't know
 		// that the requested revision is part of this channel
-		revOpts.setChannelIfUnset(snapst.TrackingChannel)
+		up.RevOpts.setChannelIfUnset(snapst.TrackingChannel)
 
 		// make sure that we switch the current channel of the snap that we're
 		// switching to
-		info.Channel = revOpts.Channel
+		info.Channel = up.RevOpts.Channel
 
 		plan.targets = append(plan.targets, target{
 			info:   info,
 			snapst: *snapst,
 			setup: SnapSetup{
-				Channel:   revOpts.Channel,
-				CohortKey: revOpts.CohortKey,
+				Channel:   up.RevOpts.Channel,
+				CohortKey: up.RevOpts.CohortKey,
 				SnapPath:  info.MountFile(),
 
 				// if the caller specified a revision, then we always run
 				// through the entire update process. this enables something
 				// like "snap refresh --revision=n", where revision n is already
 				// installed
-				AlwaysUpdate: !revOpts.Revision.Unset(),
+				AlwaysUpdate: !up.RevOpts.Revision.Unset(),
 			},
 			components: compsups,
 		})
@@ -710,8 +723,8 @@ func collectCurrentSnapsAndActions(
 	opts Options,
 	enforcedSets func() (*snapasserts.ValidationSets, error),
 	fallbackID int,
-) (actionsByUserID map[int][]*store.SnapAction, hasLocalRevision map[*SnapState]RevisionOptions, current []*store.CurrentSnap, err error) {
-	hasLocalRevision = make(map[*SnapState]RevisionOptions)
+) (actionsByUserID map[int][]*store.SnapAction, hasLocalRevision map[string]*SnapState, current []*store.CurrentSnap, err error) {
+	hasLocalRevision = make(map[string]*SnapState)
 	actionsByUserID = make(map[int][]*store.SnapAction)
 	refreshAll := len(requested) == 0
 
@@ -736,7 +749,7 @@ func collectCurrentSnapsAndActions(
 		}
 
 		if !req.RevOpts.Revision.Unset() && snapst.LastIndex(req.RevOpts.Revision) != -1 {
-			hasLocalRevision[snapst] = req.RevOpts
+			hasLocalRevision[snapst.InstanceName()] = snapst
 			return nil
 		}
 
@@ -770,7 +783,7 @@ func collectCurrentSnapsAndActions(
 		// consider this snap for a store update, but we still should return it
 		// as a target for potentially switching channels or cohort keys
 		if !action.Revision.Unset() && action.Revision == installed.Revision {
-			hasLocalRevision[snapst] = req.RevOpts
+			hasLocalRevision[installed.InstanceName] = snapst
 			return nil
 		}
 

overlord/snapstate/target.go

@@ -144,7 +144,7 @@ func (t *target) setups(st *state.State, opts Options) (SnapSetup, []ComponentSe
 		AlwaysUpdate: t.setup.AlwaysUpdate,
 
 		Base:               t.info.Base,
-		Prereq:             getKeys(providerContentAttrs),
+		Prereq:             keys(providerContentAttrs),
 		PrereqContentAttrs: providerContentAttrs,
 		UserID:             snapUserID,
 		Flags:              flags.ForSnapSetup(),
@@ -1073,6 +1073,9 @@ type StoreUpdate struct {
 	InstanceName string
 	// RevOpts contains options that apply to the update of this snap.
 	RevOpts RevisionOptions
+	// AdditionalComponents is a list of additional components to install during
+	// the refresh.
+	AdditionalComponents []string
 }
 
 // StoreUpdateGoal creates a new UpdateGoal to update snaps from the store.
@@ -1135,6 +1138,13 @@ func validateAndInitStoreUpdates(allSnaps map[string]*SnapState, updates map[str
 			return err
 		}
 
+		snapName, _ := snap.SplitInstanceName(sn.InstanceName)
+		for _, add := range sn.AdditionalComponents {
+			if snapst.CurrentComponentState(naming.NewComponentRef(snapName, add)) != nil {
+				return snap.AlreadyInstalledComponentError{Component: add}
+			}
+		}
+
 		updates[sn.InstanceName] = sn
 	}
 

overlord/snapstate/target_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/snapcore/snapd/snap/naming"
 	"github.com/snapcore/snapd/snap/snaptest"
 	"github.com/snapcore/snapd/store"
+	"github.com/snapcore/snapd/testutil"
 	. "gopkg.in/check.v1"
 )
 
@@ -382,6 +383,41 @@ func (s *targetTestSuite) TestUpdateSnapNotInstalled(c *C) {
 	c.Assert(err, ErrorMatches, `snap "some-snap" is not installed`)
 }
 
+func (s *targetTestSuite) TestUpdateAdditionalComponentAlreadyInstalled(c *C) {
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	seq := snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{{
+		RealName: "snap-1",
+		SnapID:   snaptest.AssertedSnapID("snap-1"),
+		Revision: snap.R(7),
+	}})
+
+	seq.AddComponentForRevision(snap.R(7), &sequence.ComponentState{
+		SideInfo: &snap.ComponentSideInfo{
+			Component: naming.NewComponentRef("snap-1", "comp-1"),
+			Revision:  snap.R(1),
+		},
+		CompType: snap.StandardComponent,
+	})
+
+	snapstate.Set(s.state, "snap-1", &snapstate.SnapState{
+		Active:          true,
+		TrackingChannel: "latest/stable",
+		Sequence:        seq,
+		Current:         snap.R(7),
+		SnapType:        "app",
+	})
+
+	goal := snapstate.StoreUpdateGoal(snapstate.StoreUpdate{
+		InstanceName:         "snap-1",
+		AdditionalComponents: []string{"comp-1"},
+	})
+
+	_, err := snapstate.UpdateOne(context.Background(), s.state, goal, nil, snapstate.Options{})
+	c.Assert(err, testutil.ErrorIs, snap.AlreadyInstalledComponentError{Component: "comp-1"})
+}
+
 func (s *targetTestSuite) TestInvalidPathGoals(c *C) {
 	s.state.Lock()
 	defer s.state.Unlock()