o/snapstate: add tri-state option (nil = default reservation, otherwise res_size)

Author: Rnfudge02

overlord/snapstate/snapstate.go

@@ -174,23 +174,25 @@ func ShouldSendNotificationsToTheUser(st *state.State) (bool, error) {
 	return true, nil
 }
 
-func diskSpaceReservation(tr *config.Transaction) uint64 {
+func diskSpaceReservation(tr *config.Transaction) *uint64 {
 	var reservation string
 	err := tr.Get("core", "system.disk-space-reservation", &reservation)
 	if config.IsNoOption(err) {
 		// TODO: decide how unset system.disk-space-reservation should behave when
 		// the experimental disk space feature flags are graduated.
-		return defaultDiskSpaceReservation
+		defaultReservation := defaultDiskSpaceReservation
+		return &defaultReservation
 	}
 	if err != nil {
-		return 0
+		return nil
 	}
 
 	parsedReservation, err := strutil.ParseByteSize(reservation)
 	if err != nil {
-		return 0
+		return nil
 	}
-	return uint64(parsedReservation)
+	reservationBytes := uint64(parsedReservation)
+	return &reservationBytes
 }
 
 // ConfigureSnap returns a set of tasks to configure snapName as done during installation/refresh.
@@ -2524,7 +2526,7 @@ func autoRefreshPhase2(st *state.State, candidates []*refreshCandidate, flags *F
 
 func checkDiskSpaceDownload(st *state.State, infos []minimalInstallInfo, rootDir string) error {
 	reservation := diskSpaceReservation(config.NewTransaction(st))
-	if reservation == 0 {
+	if reservation == nil {
 		return nil
 	}
 
@@ -2533,7 +2535,7 @@ func checkDiskSpaceDownload(st *state.State, infos []minimalInstallInfo, rootDir
 		totalSize += uint64(info.DownloadSize())
 	}
 
-	return checkForAvailableSpace(totalSize, reservation, infos, "download", rootDir)
+	return checkForAvailableSpace(totalSize, *reservation, infos, "download", rootDir)
 }
 
 // checkDiskSpace checks if there is enough space for the requested snaps and their prerequisites
@@ -2560,7 +2562,7 @@ func checkDiskSpace(st *state.State, changeKind string, infos []minimalInstallIn
 	}
 
 	reservation := diskSpaceReservation(tr)
-	if reservation == 0 {
+	if reservation == nil {
 		return nil
 	}
 
@@ -2569,7 +2571,7 @@ func checkDiskSpace(st *state.State, changeKind string, infos []minimalInstallIn
 		return err
 	}
 
-	return checkForAvailableSpace(totalSize, reservation, infos, changeKind, dirs.SnapdStateDir(dirs.GlobalRootDir))
+	return checkForAvailableSpace(totalSize, *reservation, infos, changeKind, dirs.SnapdStateDir(dirs.GlobalRootDir))
 }
 
 func checkForAvailableSpace(totalSize, reservation uint64, infos []minimalInstallInfo, changeKind string, rootDir string) error {
@@ -3171,11 +3173,11 @@ func Remove(st *state.State, name string, revision snap.Revision, flags *RemoveF
 	// will only be greater than 0 if the feature is enabled.
 	if snapshotSize > 0 {
 		reservation := diskSpaceReservation(config.NewTransaction(st))
-		if reservation == 0 {
+		if reservation == nil {
 			return ts, err
 		}
 
-		requiredSpace := snapshotSize + reservation
+		requiredSpace := snapshotSize + *reservation
 		path := dirs.SnapdStateDir(dirs.GlobalRootDir)
 		if err := osutilCheckFreeSpace(path, requiredSpace); err != nil {
 			if _, ok := err.(*osutil.NotEnoughDiskSpaceError); ok {
@@ -3327,7 +3329,7 @@ func removeTasks(st *state.State, snapst *SnapState, removals map[string]bool, r
 				if err != nil && !config.IsNoOption(err) {
 					return nil, 0, err
 				}
-				if checkDiskSpaceRemove && diskSpaceReservation(tr) != 0 {
+				if checkDiskSpaceRemove && diskSpaceReservation(tr) != nil {
 					snapshotSize, err = EstimateSnapshotSize(st, instanceName, nil)
 					if err != nil {
 						return nil, 0, err
@@ -3611,11 +3613,11 @@ func RemoveMany(st *state.State, names []string, flags *RemoveFlags) ([]string,
 	// will only be greater than 0 if the feature is enabled.
 	if totalSnapshotsSize > 0 {
 		reservation := diskSpaceReservation(config.NewTransaction(st))
-		if reservation == 0 {
+		if reservation == nil {
 			return removed, tasksets, nil
 		}
 
-		requiredSpace := totalSnapshotsSize + reservation
+		requiredSpace := totalSnapshotsSize + *reservation
 		if err := osutilCheckFreeSpace(path, requiredSpace); err != nil {
 			if _, ok := err.(*osutil.NotEnoughDiskSpaceError); ok {
 				return nil, nil, &InsufficientSpaceError{

overlord/snapstate/snapstate_remove_test.go

@@ -1734,17 +1734,21 @@ func (s *snapmgrTestSuite) TestRemoveManyDiskSpaceCheckPasses(c *C) {
 	c.Check(err, IsNil)
 }
 
-func (s *snapmgrTestSuite) TestRemoveManyDiskSpaceReservationZeroSkipsSnapshotSizeCheck(c *C) {
+func (s *snapmgrTestSuite) TestRemoveManyDiskSpaceReservationZeroChecksSnapshotSize(c *C) {
 	s.state.Lock()
 	defer s.state.Unlock()
 
+	var snapshotSizeCall int
 	snapstate.EstimateSnapshotSize = func(st *state.State, instanceName string, users []string) (uint64, error) {
-		c.Fatalf("unexpected snapshot size estimation")
-		return 0, nil
+		snapshotSizeCall++
+		c.Check(instanceName, Equals, "one")
+		return 123, nil
 	}
 
+	var requiredSizes []uint64
 	restore := snapstate.MockOsutilCheckFreeSpace(func(path string, required uint64) error {
-		c.Fatalf("unexpected disk space check")
+		c.Check(path, Equals, filepath.Join(dirs.GlobalRootDir, "/var/lib/snapd"))
+		requiredSizes = append(requiredSizes, required)
 		return nil
 	})
 	defer restore()
@@ -1773,6 +1777,9 @@ func (s *snapmgrTestSuite) TestRemoveManyDiskSpaceReservationZeroSkipsSnapshotSi
 	_, _, err := snapstate.RemoveMany(s.state, []string{"one"}, nil)
 	c.Assert(err, IsNil)
 	c.Check(automaticSnapshotCalled, Equals, true)
+	// 0B reservation means checks run with just the snapshot size, no buffer
+	c.Check(snapshotSizeCall, Equals, 1)
+	c.Check(requiredSizes, DeepEquals, []uint64{123})
 }
 
 type snapdBackend struct {

overlord/snapstate/snapstate_update_test.go

@@ -7473,6 +7473,16 @@ func (s *snapmgrTestSuite) TestUpdateDiskSpaceDefaultReservationError(c *C) {
 	c.Check(diskSpaceErr.Snaps, DeepEquals, []string{"some-snap"})
 }
 
+func (s *snapmgrTestSuite) TestUpdateDiskSpaceDefaultReservationHappy(c *C) {
+	featureFlag := true
+	failInstallSize := false
+	failDiskCheck := false
+	// No system.disk-space-reservation is set, so the default 5MB reservation is used.
+	// The helper asserts required size == 123 + DefaultDiskSpaceReservation.
+	err := s.testUpdateDiskSpaceCheck(c, featureFlag, failInstallSize, failDiskCheck)
+	c.Check(err, IsNil)
+}
+
 func (s *snapmgrTestSuite) TestUpdateConfigureDiskSpaceReservation(c *C) {
 	const freeDiskSpace = uint64(1500)
 	var requiredSizes []uint64
@@ -7521,16 +7531,21 @@ func (s *snapmgrTestSuite) TestUpdateConfigureDiskSpaceReservation(c *C) {
 	c.Check(requiredSizes, DeepEquals, []uint64{2123, 1123})
 }
 
-func (s *snapmgrTestSuite) TestUpdateDiskSpaceReservationZeroSkipsCheck(c *C) {
+func (s *snapmgrTestSuite) TestUpdateDiskSpaceReservationZeroChecksNormalSize(c *C) {
+	var requiredSizes []uint64
 	restore := snapstate.MockOsutilCheckFreeSpace(func(path string, sz uint64) error {
-		c.Fatalf("unexpected disk space check")
+		c.Check(path, Equals, filepath.Join(dirs.GlobalRootDir, "/var/lib/snapd"))
+		requiredSizes = append(requiredSizes, sz)
 		return nil
 	})
 	defer restore()
 
+	var installSizeCalled bool
 	restore = snapstate.MockInstallSize(func(st *state.State, snaps []snapstate.MinimalInstallInfo, userID int, prqt snapstate.PrereqTracker) (uint64, error) {
-		c.Fatalf("unexpected install size calculation")
-		return 0, nil
+		installSizeCalled = true
+		c.Assert(snaps, HasLen, 1)
+		c.Check(snaps[0].InstanceName(), Equals, "some-snap")
+		return 123, nil
 	})
 	defer restore()
 
@@ -7554,6 +7569,9 @@ func (s *snapmgrTestSuite) TestUpdateDiskSpaceReservationZeroSkipsCheck(c *C) {
 	opts := &snapstate.RevisionOptions{Channel: "some-channel"}
 	_, err := snapstate.Update(s.state, "some-snap", opts, s.user.ID, snapstate.Flags{})
 	c.Assert(err, IsNil)
+	// 0B reservation means checks run with just the install size, no buffer
+	c.Check(installSizeCalled, Equals, true)
+	c.Check(requiredSizes, DeepEquals, []uint64{123})
 }
 
 func (s *snapmgrTestSuite) TestUpdateDiskCheckSkippedIfDisabled(c *C) {