Handle refresh hints (WIP).refresh-control/refresh-hints

11 files changed, 295 insertions, 83 deletions

diff --git a/overlord/snapstate/autorefresh.go b/overlord/snapstate/autorefresh.go
index 16c72cce7c..7a19562abd 100644
--- a/overlord/snapstate/autorefresh.go
+++ b/overlord/snapstate/autorefresh.go
@@ -58,6 +58,38 @@ var (
 // refreshRetryDelay specified the minimum time to retry failed refreshes
 var refreshRetryDelay = 20 * time.Minute
 
+// refreshCandidate carries information about a single snap to update as a part
+// of auto-refresh.
+type refreshCandidate struct {
+	SnapSetup
+}
+
+func (rh *refreshCandidate) Type() snap.Type {
+	return rh.SnapSetup.Type
+}
+
+func (rh *refreshCandidate) Base() string {
+	return rh.SnapSetup.Base
+}
+
+func (rh *refreshCandidate) InstanceName() string {
+	return rh.SnapSetup.InstanceName()
+}
+
+func (rh *refreshCandidate) MakeSnapSetup(st *state.State, snapst *SnapState, revnoOpts *RevisionOptions, snapUserID int) (*SnapSetup, error) {
+	rh.Channel = revnoOpts.Channel
+	rh.CohortKey = revnoOpts.CohortKey
+	return &rh.SnapSetup, nil
+}
+
+func (rh *refreshCandidate) Size() int64 {
+	return rh.DownloadInfo.Size
+}
+
+func (rh *refreshCandidate) DefaultContentPlugProviders(st *state.State) []string {
+	return rh.Prereq
+}
+
 // autoRefresh will ensure that snaps are refreshed automatically
 // according to the refresh schedule.
 type autoRefresh struct {
diff --git a/overlord/snapstate/export_test.go b/overlord/snapstate/export_test.go
index e86398013c..67041cc1cf 100644
--- a/overlord/snapstate/export_test.go
+++ b/overlord/snapstate/export_test.go
@@ -233,6 +233,10 @@ var (
 
 type AuxStoreInfo = auxStoreInfo
 
+type ManualUpdateInfo = manualUpdateInfo
+
+var NewManualUpdateInfo = newManualUpdateInfo
+
 // link, misc handlers
 var (
	MissingDisabledServices = missingDisabledServices
@@ -258,7 +262,7 @@ func MockCurrentSnaps(f func(st *state.State) ([]*store.CurrentSnap, error)) fun
	}
 }
 
-func MockInstallSize(f func(st *state.State, snaps []*snap.Info, userID int) (uint64, error)) func() {
+func MockInstallSize(f func(st *state.State, snaps []UpdateInfo, userID int) (uint64, error)) func() {
	old := installSize
	installSize = f
	return func() {
diff --git a/overlord/snapstate/handlers_prereq_test.go b/overlord/snapstate/handlers_prereq_test.go
index 397c4e310c..44b934430f 100644
--- a/overlord/snapstate/handlers_prereq_test.go
+++ b/overlord/snapstate/handlers_prereq_test.go
@@ -61,7 +61,7 @@ func (s *prereqSuite) SetUpTest(c *C) {
	restoreCheckFreeSpace := snapstate.MockOsutilCheckFreeSpace(func(string, uint64) error { return nil })
	s.AddCleanup(restoreCheckFreeSpace)
 
-	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	return 0, nil
	})
	s.AddCleanup(restoreInstallSize)
diff --git a/overlord/snapstate/refreshhints.go b/overlord/snapstate/refreshhints.go
index 537266a8e1..e81d5a7b72 100644
--- a/overlord/snapstate/refreshhints.go
+++ b/overlord/snapstate/refreshhints.go
@@ -22,9 +22,11 @@ package snapstate
 import (
	"time"
 
+	"github.com/snapcore/snapd/logger"
	"github.com/snapcore/snapd/overlord/auth"
	"github.com/snapcore/snapd/overlord/state"
	"github.com/snapcore/snapd/release"
+	"github.com/snapcore/snapd/snap"
	"github.com/snapcore/snapd/store"
	"github.com/snapcore/snapd/timings"
 )
@@ -71,12 +73,24 @@ func (r *refreshHints) refresh() error {
	perfTimings := timings.New(map[string]string{"ensure": "refresh-hints"})
	defer perfTimings.Save(r.state)
 
+	var updates []*snap.Info
+	var ignoreValidationByInstanceName map[string]bool
	timings.Run(perfTimings, "refresh-candidates", "query store for refresh candidates", func(tm timings.Measurer) {
-	_, _, _, err = refreshCandidates(auth.EnsureContextTODO(), r.state, nil, nil, &store.RefreshOptions{RefreshManaged: refreshManaged})
+	updates, _, ignoreValidationByInstanceName, err = refreshCandidates(auth.EnsureContextTODO(), r.state, nil, nil, &store.RefreshOptions{RefreshManaged: refreshManaged})
	})
	// TODO: we currently set last-refresh-hints even when there was an
	// error. In the future we may retry with a backoff.
	r.state.Set("last-refresh-hints", time.Now())
+	if err == nil {
+	deviceCtx, err := DeviceCtxFromState(r.state, nil)
+	if err != nil {
+	return err
+	}
+	hints, err := refreshHintsFromCandidates(r.state, updates, ignoreValidationByInstanceName, deviceCtx)
+	if err == nil {
+	r.state.Set("refresh-candidates", hints)
+	}
+	}
	return err
 }
 
@@ -121,3 +135,79 @@ func (r *refreshHints) Ensure() error {
	}
	return r.refresh()
 }
+
+func refreshHintsFromCandidates(st *state.State, updates []*snap.Info, ignoreValidationByInstanceName map[string]bool, deviceCtx DeviceContext) ([]*refreshCandidate, error) {
+	if ValidateRefreshes != nil && len(updates) != 0 {
+	userID := 0
+	var err error
+	updates, err = ValidateRefreshes(st, updates, ignoreValidationByInstanceName, userID, deviceCtx)
+	if err != nil {
+	return nil, err
+	}
+	}
+
+	hints := make([]*refreshCandidate, len(updates))
+	for _, update := range updates {
+	var snapst SnapState
+	if err := Get(st, update.InstanceName(), &snapst); err != nil {
+	return nil, err
+	}
+
+	flags := snapst.Flags
+	flags.IsAutoRefresh = true
+	flags, err := earlyChecks(st, &snapst, update, flags)
+	if err != nil {
+	logger.Noticef("cannot update %q: %v", update.InstanceName(), err)
+	continue
+	}
+
+	if err := earlyEpochCheck(update, &snapst); err != nil {
+	logger.Noticef("cannot update %q: %v", update.InstanceName(), err)
+	continue
+	}
+	snapsup := &refreshCandidate{
+	SnapSetup{
+	Base: update.Base,
+	Prereq: defaultContentPlugProviders(st, update),
+	DownloadInfo: &update.DownloadInfo,
+	SideInfo: &update.SideInfo,
+	Type: update.Type(),
+	PlugsOnly: len(update.Slots) == 0,
+	Flags: flags.ForSnapSetup(),
+	InstanceKey: update.InstanceKey,
+	auxStoreInfo: auxStoreInfo{
+	Website: update.Website,
+	Media: update.Media,
+	},
+	}}
+
+	hints = append(hints, snapsup)
+	}
+	return hints, nil
+}
+
+func updatesFromRefreshHints(st *state.State) ([]*refreshCandidate, map[string]*SnapState, error) {
+	var hints []*refreshCandidate
+	err := st.Get("refresh-candidates", &hints)
+	if err != nil && err != state.ErrNoState {
+	return nil, nil, err
+	}
+	if err != nil {
+	return nil, nil, nil
+	}
+
+	snapStates, err := All(st)
+	if err != nil {
+	return nil, nil, err
+	}
+
+	stateByInstanceName := make(map[string]*SnapState, len(snapStates))
+	for _, hint := range hints {
+	if _, ok := snapStates[hint.SnapSetup.InstanceName()]; !ok {
+	return nil, nil, snap.NotInstalledError{Snap: hint.SnapSetup.InstanceName()}
+	}
+	}
+
+	return hints, stateByInstanceName, nil
+}
+
diff --git a/overlord/snapstate/refreshhints_test.go b/overlord/snapstate/refreshhints_test.go
index c2cdf1f6ae..089e5cba55 100644
--- a/overlord/snapstate/refreshhints_test.go
+++ b/overlord/snapstate/refreshhints_test.go
@@ -27,6 +27,7 @@ import (
 
	"github.com/snapcore/snapd/overlord/auth"
	"github.com/snapcore/snapd/overlord/snapstate"
+	"github.com/snapcore/snapd/overlord/snapstate/snapstatetest"
	"github.com/snapcore/snapd/overlord/state"
	"github.com/snapcore/snapd/release"
	"github.com/snapcore/snapd/snap"
@@ -63,6 +64,7 @@ type refreshHintsTestSuite struct {
	state *state.State
 
	store *recordingStore
+	restoreModel func()
 }
 
 var _ = Suite(&refreshHintsTestSuite{})
@@ -91,11 +93,14 @@ func (s *refreshHintsTestSuite) SetUpTest(c *C) {
	}
 
	s.state.Set("refresh-privacy-key", "privacy-key")
+
+	s.restoreModel = snapstatetest.MockDeviceModel(DefaultModel())
 }
 
 func (s *refreshHintsTestSuite) TearDownTest(c *C) {
	snapstate.CanAutoRefresh = nil
	snapstate.AutoAliases = nil
+	s.restoreModel()
 }
 
 func (s *refreshHintsTestSuite) TestLastRefresh(c *C) {
diff --git a/overlord/snapstate/snapstate.go b/overlord/snapstate/snapstate.go
index ae6d54ce54..ee68acb06c 100644
--- a/overlord/snapstate/snapstate.go
+++ b/overlord/snapstate/snapstate.go
@@ -665,7 +665,7 @@ func validateFeatureFlags(st *state.State, info *snap.Info) error {
	return nil
 }
 
-func ensureInstallPreconditions(st *state.State, info *snap.Info, flags Flags, snapst *SnapState, deviceCtx DeviceContext) (Flags, error) {
+func ensureInstallPreconditions(st *state.State, info *snap.Info, flags Flags, snapst *SnapState) (Flags, error) {
	// if snap is allowed to be devmode via the dangerous model and it's
	// confinement is indeed devmode, promote the flags.DevMode to true
	if flags.ApplySnapDevMode && info.NeedsDevMode() {
@@ -757,7 +757,7 @@ func InstallPath(st *state.State, si *snap.SideInfo, path, instanceName, channel
	}
	info.InstanceKey = instanceKey
 
-	flags, err = ensureInstallPreconditions(st, info, flags, &snapst, deviceCtx)
+	flags, err = ensureInstallPreconditions(st, info, flags, &snapst)
	if err != nil {
	return nil, nil, err
	}
@@ -844,7 +844,7 @@ func InstallWithDeviceContext(ctx context.Context, st *state.State, name string,
	return nil, fmt.Errorf("unexpected snap type %q, instead of 'base'", info.Type())
	}
 
-	flags, err = ensureInstallPreconditions(st, info, flags, &snapst, deviceCtx)
+	flags, err = ensureInstallPreconditions(st, info, flags, &snapst)
	if err != nil {
	return nil, err
	}
@@ -857,7 +857,7 @@ func InstallWithDeviceContext(ctx context.Context, st *state.State, name string,
	if checkDiskSpaceInstall {
	// check if there is enough disk space for requested snap and its
	// prerequisites.
-	totalSize, err := installSize(st, []*snap.Info{info}, userID)
+	totalSize, err := installSize(st, []UpdateInfo{newManualUpdateInfo(info, &flags)}, userID)
	if err != nil {
	return nil, err
	}
@@ -945,9 +945,9 @@ func InstallMany(st *state.State, names []string, userID int) ([]string, []*stat
	if checkDiskSpaceInstall {
	// check if there is enough disk space for requested snaps and their
	// prerequisites.
-	snapInfos := make([]*snap.Info, len(installs))
+	snapInfos := make([]UpdateInfo, len(installs))
	for i, sar := range installs {
-	snapInfos[i] = sar.Info
+	snapInfos[i] = newManualUpdateInfo(sar.Info, nil)
	}
	totalSize, err := installSize(st, snapInfos, userID)
	if err != nil {
@@ -974,7 +974,7 @@ func InstallMany(st *state.State, names []string, userID int) ([]string, []*stat
	var snapst SnapState
	var flags Flags
 
-	flags, err := ensureInstallPreconditions(st, info, flags, &snapst, deviceCtx)
+	flags, err := ensureInstallPreconditions(st, info, flags, &snapst)
	if err != nil {
	return nil, nil, err
	}
@@ -1048,6 +1048,7 @@ func updateManyFiltered(ctx context.Context, st *state.State, names []string, us
	return nil, nil, err
	}
 
+	var toUpdate []UpdateInfo
	refreshOpts := &store.RefreshOptions{IsAutoRefresh: flags.IsAutoRefresh}
	updates, stateByInstanceName, ignoreValidation, err := refreshCandidates(ctx, st, names, user, refreshOpts)
	if err != nil {
@@ -1076,14 +1077,29 @@ func updateManyFiltered(ctx context.Context, st *state.State, names []string, us
	}
	}
 
-	params := func(update *snap.Info) (*RevisionOptions, Flags, *SnapState) {
-	snapst := stateByInstanceName[update.InstanceName()]
+	refreshAll := len(names) == 0
+	for _, up := range updates {
+	snapst := stateByInstanceName[up.InstanceName()]
+	fl, err := earlyChecks(st, snapst, up, snapst.Flags)
+	if err != nil {
+	if refreshAll {
+	logger.Noticef("cannot update %q: %v", up.InstanceName(), err)
+	continue
+	}
+	return nil, nil, err
+	}
+	fl.IsAutoRefresh = flags.IsAutoRefresh
+	toUpdate = append(toUpdate, newManualUpdateInfo(up, &fl))
+	}
+
+	params := func(up UpdateInfo) (*RevisionOptions, *SnapState) {
+	snapst := stateByInstanceName[up.InstanceName()]
	// setting options to what's in state as multi-refresh doesn't let you change these
	opts := &RevisionOptions{
	Channel: snapst.TrackingChannel,
	CohortKey: snapst.CohortKey,
	}
-	return opts, snapst.Flags, snapst
+	return opts, snapst
 
	}
 
@@ -1095,15 +1111,15 @@ func updateManyFiltered(ctx context.Context, st *state.State, names []string, us
	if checkDiskSpaceRefresh {
	// check if there is enough disk space for requested snap and its
	// prerequisites.
-	totalSize, err := installSize(st, updates, userID)
+	totalSize, err := installSize(st, toUpdate, userID)
	if err != nil {
	return nil, nil, err
	}
	requiredSpace := safetyMarginDiskSpace(totalSize)
	path := dirs.SnapdStateDir(dirs.GlobalRootDir)
	if err := osutilCheckFreeSpace(path, requiredSpace); err != nil {
-	snaps := make([]string, len(updates))
-	for i, up := range updates {
+	snaps := make([]string, len(toUpdate))
+	for i, up := range toUpdate {
	snaps[i] = up.InstanceName()
	}
	if _, ok := err.(*osutil.NotEnoughDiskSpaceError); ok {
@@ -1117,15 +1133,103 @@ func updateManyFiltered(ctx context.Context, st *state.State, names []string, us
	}
	}
 
-	updated, tasksets, err := doUpdate(ctx, st, names, updates, params, userID, flags, deviceCtx, fromChange)
+	updated, tasksets, err := doUpdate(ctx, st, names, toUpdate, params, userID, flags, deviceCtx, fromChange)
	if err != nil {
	return nil, nil, err
	}
-	tasksets = finalizeUpdate(st, tasksets, len(updates) > 0, updated, userID, flags)
+	tasksets = finalizeUpdate(st, tasksets, len(toUpdate) > 0, updated, userID, flags)
	return updated, tasksets, nil
 }
 
-func doUpdate(ctx context.Context, st *state.State, names []string, updates []*snap.Info, params func(*snap.Info) (*RevisionOptions, Flags, *SnapState), userID int, globalFlags *Flags, deviceCtx DeviceContext, fromChange string) ([]string, []*state.TaskSet, error) {
+type UpdateInfo interface {
+	MakeSnapSetup(*state.State, *SnapState, *RevisionOptions, int) (*SnapSetup, error)
+	InstanceName() string
+	Type() snap.Type
+	Base() string
+	// for install size calculations
+	Size() int64
+	DefaultContentPlugProviders(*state.State) []string
+}
+
+// ByType supports sorting by snap type. The most important types come first.
+type byType []UpdateInfo
+
+func (r byType) Len() int { return len(r) }
+func (r byType) Swap(i, j int) { r[i], r[j] = r[j], r[i] }
+func (r byType) Less(i, j int) bool {
+	return r[i].Type().SortsBefore(r[j].Type())
+}
+
+type manualUpdateInfo struct {
+	snap.Info
+	flags Flags
+}
+
+func newManualUpdateInfo(info *snap.Info, flags *Flags) *manualUpdateInfo {
+	up := &manualUpdateInfo{}
+	up.Info = *info
+	if flags != nil {
+	up.flags = *flags
+	}
+	return up
+}
+
+func (up *manualUpdateInfo) InstanceName() string {
+	return up.Info.InstanceName()
+}
+
+func (up *manualUpdateInfo) Type() snap.Type {
+	return up.SnapType
+}
+
+func (up *manualUpdateInfo) Base() string {
+	return up.Info.Base
+}
+
+func (up *manualUpdateInfo) Size() int64 {
+	return up.DownloadInfo.Size
+}
+
+func (up *manualUpdateInfo) DefaultContentPlugProviders(st *state.State) []string {
+	return defaultContentPlugProviders(st, &up.Info)
+}
+
+func earlyChecks(st *state.State, snapst *SnapState, update *snap.Info, flags Flags) (Flags, error) {
+	var err error
+	flags, err = ensureInstallPreconditions(st, update, flags, snapst)
+	if err != nil {
+	return flags, err
+	}
+
+	if err := earlyEpochCheck(update, snapst); err != nil {
+	return flags, err
+	}
+	return flags, nil
+}
+
+func (up *manualUpdateInfo) MakeSnapSetup(st *state.State, snapst *SnapState, revnoOpts *RevisionOptions, snapUserID int) (*SnapSetup, error) {
+	inf := &up.Info
+	snapsup := &SnapSetup{
+	Base: up.Info.Base,
+	Prereq: defaultContentPlugProviders(st, inf),
+	Channel: revnoOpts.Channel,
+	CohortKey: revnoOpts.CohortKey,
+	UserID: snapUserID,
+	Flags: up.flags.ForSnapSetup(),
+	DownloadInfo: &inf.DownloadInfo,
+	SideInfo: &inf.SideInfo,
+	Type: inf.Type(),
+	PlugsOnly: len(inf.Slots) == 0,
+	InstanceKey: inf.InstanceKey,
+	auxStoreInfo: auxStoreInfo{
+	Website: inf.Website,
+	Media: inf.Media,
+	},
+	}
+	return snapsup, nil
+}
+
+func doUpdate(ctx context.Context, st *state.State, names []string, updates []UpdateInfo, params func(UpdateInfo) (*RevisionOptions, *SnapState), userID int, globalFlags *Flags, deviceCtx DeviceContext, fromChange string) ([]string, []*state.TaskSet, error) {
	if globalFlags == nil {
	globalFlags = &Flags{}
	}
@@ -1171,7 +1275,7 @@ func doUpdate(ctx context.Context, st *state.State, names []string, updates []*s
	}
 
	// first snapd, core, bases, then rest
-	sort.Stable(snap.ByType(updates))
+	sort.Stable(byType(updates))
	prereqs := make(map[string]*state.TaskSet)
	waitPrereq := func(ts *state.TaskSet, prereqName string) {
	preTs := prereqs[prereqName]
@@ -1184,19 +1288,14 @@ func doUpdate(ctx context.Context, st *state.State, names []string, updates []*s
	// updates is sorted by kind so this will process first core
	// and bases and then other snaps
	for _, update := range updates {
-	revnoOpts, flags, snapst := params(update)
-	flags.IsAutoRefresh = globalFlags.IsAutoRefresh
-
-	flags, err := ensureInstallPreconditions(st, update, flags, snapst, deviceCtx)
+	revnoOpts, snapst := params(update)
+	snapUserID, err := userIDForSnap(st, snapst, userID)
	if err != nil {
-	if refreshAll {
-	logger.Noticef("cannot update %q: %v", update.InstanceName(), err)
-	continue
-	}
	return nil, nil, err
	}
 
-	if err := earlyEpochCheck(update, snapst); err != nil {
+	snapsup, err := update.MakeSnapSetup(st, snapst, revnoOpts, snapUserID)
+	if err != nil {
	if refreshAll {
	logger.Noticef("cannot update %q: %v", update.InstanceName(), err)
	continue
@@ -1204,29 +1303,6 @@ func doUpdate(ctx context.Context, st *state.State, names []string, updates []*s
	return nil, nil, err
	}
 
-	snapUserID, err := userIDForSnap(st, snapst, userID)
-	if err != nil {
-	return nil, nil, err
-	}
-
-	snapsup := &SnapSetup{
-	Base: update.Base,
-	Prereq: defaultContentPlugProviders(st, update),
-	Channel: revnoOpts.Channel,
-	CohortKey: revnoOpts.CohortKey,
-	UserID: snapUserID,
-	Flags: flags.ForSnapSetup(),
-	DownloadInfo: &update.DownloadInfo,
-	SideInfo: &update.SideInfo,
-	Type: update.Type(),
-	PlugsOnly: len(update.Slots) == 0,
-	InstanceKey: update.InstanceKey,
-	auxStoreInfo: auxStoreInfo{
-	Website: update.Website,
-	Media: update.Media,
-	},
-	}
-
	ts, err := doInstall(st, snapst, snapsup, 0, fromChange, inUseFor(deviceCtx))
	if err != nil {
	if refreshAll {
@@ -1252,8 +1328,8 @@ func doUpdate(ctx context.Context, st *state.State, names []string, updates []*s
	// snaps
	waitPrereq(ts, defaultCoreSnapName)
	waitPrereq(ts, "snapd")
-	if update.Base != "" {
-	waitPrereq(ts, update.Base)
+	if update.Base() != "" {
+	waitPrereq(ts, update.Base())
	}
	}
	// keep track of kernel/gadget udpates
@@ -1343,7 +1419,7 @@ func applyAutoAliasesDelta(st *state.State, delta map[string][]string, op string
	return applyTs, nil
 }
 
-func autoAliasesUpdate(st *state.State, names []string, updates []*snap.Info) (changed map[string][]string, mustPrune map[string][]string, transferTargets map[string]bool, err error) {
+func autoAliasesUpdate(st *state.State, names []string, updates []UpdateInfo) (changed map[string][]string, mustPrune map[string][]string, transferTargets map[string]bool, err error) {
	changed, dropped, err := autoAliasesDelta(st, nil)
	if err != nil {
	if len(names) != 0 {
@@ -1652,11 +1728,15 @@ func UpdateWithDeviceContext(st *state.State, name string, opts *RevisionOptions
	flags.Classic = flags.Classic || snapst.Flags.Classic
	}
 
-	var updates []*snap.Info
+	var updates []UpdateInfo
	info, infoErr := infoForUpdate(st, &snapst, name, opts, userID, flags, deviceCtx)
	switch infoErr {
	case nil:
-	updates = append(updates, info)
+	flags, err = earlyChecks(st, &snapst, info, flags)
+	if err != nil {
+	return nil, err
+	}
+	updates = append(updates, newManualUpdateInfo(info, &flags))
	case store.ErrNoUpdateAvailable:
	// there may be some new auto-aliases
	default:
@@ -1668,6 +1748,7 @@ func UpdateWithDeviceContext(st *state.State, name string, opts *RevisionOptions
	if err != nil && !config.IsNoOption(err) {
	return nil, err
	}
+
	if checkDiskSpaceRefresh {
	// check if there is enough disk space for requested snap and its
	// prerequisites.
@@ -1693,8 +1774,8 @@ func UpdateWithDeviceContext(st *state.State, name string, opts *RevisionOptions
	}
	}
 
-	params := func(update *snap.Info) (*RevisionOptions, Flags, *SnapState) {
-	return opts, flags, &snapst
+	params := func(UpdateInfo) (*RevisionOptions, *SnapState) {
+	return opts, &snapst
	}
 
	_, tts, err := doUpdate(context.TODO(), st, []string{name}, updates, params, userID, &flags, deviceCtx, fromChange)
diff --git a/overlord/snapstate/snapstate_install_test.go b/overlord/snapstate/snapstate_install_test.go
index 5186650f38..fca974b908 100644
--- a/overlord/snapstate/snapstate_install_test.go
+++ b/overlord/snapstate/snapstate_install_test.go
@@ -1773,7 +1773,7 @@ func (s *snapmgrTestSuite) TestInstallFirstLocalRunThrough(c *C) {
	// use the real thing for this one
	snapstate.MockOpenSnapFile(backend.OpenSnapFile)
 
-	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	c.Fatalf("installSize shouldn't be hit with local install")
	return 0, nil
	})
@@ -2897,7 +2897,7 @@ func (s *snapmgrTestSuite) TestInstallDiskSpaceError(c *C) {
 }
 
 func (s *snapmgrTestSuite) TestInstallSizeError(c *C) {
-	restore := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restore := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	return 0, fmt.Errorf("boom")
	})
	defer restore()
diff --git a/overlord/snapstate/snapstate_test.go b/overlord/snapstate/snapstate_test.go
index e1d6f5b8b7..314b27196b 100644
--- a/overlord/snapstate/snapstate_test.go
+++ b/overlord/snapstate/snapstate_test.go
@@ -168,7 +168,7 @@ func (s *snapmgrTestSuite) SetUpTest(c *C) {
	snapstate.EstimateSnapshotSize = func(st *state.State, instanceName string, users []string) (uint64, error) {
	return 1, nil
	}
-	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	return 0, nil
	})
	s.AddCleanup(restoreInstallSize)
diff --git a/overlord/snapstate/snapstate_update_test.go b/overlord/snapstate/snapstate_update_test.go
index b37681c151..780697f0ca 100644
--- a/overlord/snapstate/snapstate_update_test.go
+++ b/overlord/snapstate/snapstate_update_test.go
@@ -5003,7 +5003,7 @@ func (s *snapmgrTestSuite) testUpdateManyDiskSpaceCheck(c *C, featureFlag, failD
	})
	defer restore()
 
-	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	installSizeCalled = true
	if failInstallSize {
	return 0, fmt.Errorf("boom")
@@ -5628,7 +5628,7 @@ func (s *snapmgrTestSuite) testUpdateDiskSpaceCheck(c *C, featureFlag, failInsta
 
	var installSizeCalled bool
 
-	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+	restoreInstallSize := snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.UpdateInfo, userID int) (uint64, error) {
	installSizeCalled = true
	if failInstallSize {
	return 0, fmt.Errorf("boom")
diff --git a/overlord/snapstate/storehelpers.go b/overlord/snapstate/storehelpers.go
index 03d3711aa3..c8a3f51b31 100644
--- a/overlord/snapstate/storehelpers.go
+++ b/overlord/snapstate/storehelpers.go
@@ -89,7 +89,7 @@ func refreshOptions(st *state.State, origOpts *store.RefreshOptions) (*store.Ref
 // potentially more than once. It assumes the initial list of snaps already has
 // download infos set.
 // The state must be locked by the caller.
-var installSize = func(st *state.State, snaps []*snap.Info, userID int) (uint64, error) {
+var installSize = func(st *state.State, snaps []UpdateInfo, userID int) (uint64, error) {
	curSnaps, err := currentSnaps(st)
	if err != nil {
	return 0, err
@@ -107,21 +107,21 @@ var installSize = func(st *state.State, snaps []*snap.Info, userID int) (uint64,
 
	var prereqs []string
 
-	resolveBaseAndContentProviders := func(snapInfo *snap.Info) {
-	if snapInfo.SnapType != snap.TypeApp {
+	resolveBaseAndContentProviders := func(inst UpdateInfo) {
+	if inst.Type() != snap.TypeApp {
	return
	}
-	if snapInfo.Base != "none" {
+	if inst.Base() != "none" {
	base := defaultCoreSnapName
-	if snapInfo.Base != "" {
-	base = snapInfo.Base
+	if inst.Base() != "" {
+	base = inst.Base()
	}
	if !accountedSnaps[base] {
	prereqs = append(prereqs, base)
	accountedSnaps[base] = true
	}
	}
-	for _, snapName := range defaultContentPlugProviders(st, snapInfo) {
+	for _, snapName := range inst.DefaultContentPlugProviders(st) {
	if !accountedSnaps[snapName] {
	prereqs = append(prereqs, snapName)
	accountedSnaps[snapName] = true
@@ -130,12 +130,12 @@ var installSize = func(st *state.State, snaps []*snap.Info, userID int) (uint64,
	}
 
	snapSizes := map[string]uint64{}
-	for _, snapInfo := range snaps {
-	if snapInfo.DownloadInfo.Size == 0 {
-	return 0, fmt.Errorf("internal error: download info missing for %q", snapInfo.InstanceName())
+	for _, inst := range snaps {
+	if inst.Size() == 0 {
+	return 0, fmt.Errorf("internal error: download info missing for %q", inst.InstanceName())
	}
-	snapSizes[snapInfo.InstanceName()] = uint64(snapInfo.Size)
-	resolveBaseAndContentProviders(snapInfo)
+	snapSizes[inst.InstanceName()] = uint64(inst.Size())
+	resolveBaseAndContentProviders(inst)
	}
 
	opts, err := refreshOptions(st, nil)
@@ -170,7 +170,7 @@ var installSize = func(st *state.State, snaps []*snap.Info, userID int) (uint64,
	for _, res := range results {
	snapSizes[res.InstanceName()] = uint64(res.Size)
	// results may have new base or content providers
-	resolveBaseAndContentProviders(res.Info)
+	resolveBaseAndContentProviders(newManualUpdateInfo(res.Info, nil))
	}
	}
 
diff --git a/overlord/snapstate/storehelpers_test.go b/overlord/snapstate/storehelpers_test.go
index f6494a9fd6..fbfe23b0f3 100644
--- a/overlord/snapstate/storehelpers_test.go
+++ b/overlord/snapstate/storehelpers_test.go
@@ -190,7 +190,7 @@ func (s *snapmgrTestSuite) TestInstallSizeSimple(c *C) {
	})
	snap2.Size = snap2Size
 
-	sz, err := snapstate.InstallSize(st, []*snap.Info{snap1, snap2}, 0)
+	sz, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil), snapstate.NewManualUpdateInfo(snap2, nil)}, 0)
	c.Assert(err, IsNil)
	c.Check(sz, Equals, uint64(snap1Size+snap2Size))
 }
@@ -232,7 +232,7 @@ func (s *snapmgrTestSuite) TestInstallSizeWithBases(c *C) {
	Current: snap.R(1),
	})
 
-	sz, err := snapstate.InstallSize(st, []*snap.Info{snap1, snap2, snap3, snap4}, 0)
+	sz, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil), snapstate.NewManualUpdateInfo(snap2, nil), snapstate.NewManualUpdateInfo(snap3, nil), snapstate.NewManualUpdateInfo(snap4, nil)}, 0)
	c.Assert(err, IsNil)
	c.Check(sz, Equals, uint64(snap1Size+snap2Size+snap3Size+snap4Size+someBaseSize+otherBaseSize))
 }
@@ -262,7 +262,7 @@ func (s *snapmgrTestSuite) TestInstallSizeWithContentProviders(c *C) {
	s.mockCoreSnap(c)
 
	// both snaps have same content providers and base
-	sz, err := snapstate.InstallSize(st, []*snap.Info{snap1, snap2}, 0)
+	sz, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil), snapstate.NewManualUpdateInfo(snap2, nil)}, 0)
	c.Assert(err, IsNil)
	c.Check(sz, Equals, uint64(snap1Size+snap2Size+someBaseSize+snapContentSlotSize))
 }
@@ -284,7 +284,7 @@ func (s *snapmgrTestSuite) TestInstallSizeWithNestedDependencies(c *C) {
 
	s.mockCoreSnap(c)
 
-	sz, err := snapstate.InstallSize(st, []*snap.Info{snap1}, 0)
+	sz, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil)}, 0)
	c.Assert(err, IsNil)
	c.Check(sz, Equals, uint64(snap1Size+someBaseSize+snapOtherContentSlotSize+someOtherBaseSize))
 }
@@ -323,7 +323,7 @@ func (s *snapmgrTestSuite) TestInstallSizeWithOtherChangeAffectingSameSnaps(c *C
	})
	snap3.Size = snap3Size
 
-	sz, err := snapstate.InstallSize(st, []*snap.Info{snap1, snap3}, 0)
+	sz, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil), snapstate.NewManualUpdateInfo(snap3, nil)}, 0)
	c.Assert(err, IsNil)
	// snap3 and its base installed by another change, not counted here
	c.Check(sz, Equals, uint64(snap1Size+someBaseSize))
@@ -342,6 +342,6 @@ func (s *snapmgrTestSuite) TestInstallSizeErrorNoDownloadInfo(c *C) {
	RealName: "snap",
	}}
 
-	_, err := snapstate.InstallSize(st, []*snap.Info{snap1}, 0)
+	_, err := snapstate.InstallSize(st, []snapstate.UpdateInfo{snapstate.NewManualUpdateInfo(snap1, nil)}, 0)
	c.Assert(err, ErrorMatches, `internal error: download info missing.*`)
 }