diff --git a/pkg/action/upgrade.go b/pkg/action/upgrade.go
index e5c92e2..59c1104 100644
--- a/pkg/action/upgrade.go
+++ b/pkg/action/upgrade.go
@@ -153,6 +153,9 @@ func (u *UpgradeAction) Run() (err error) {
 	cleanup.Push(func() error { return u.remove(upgradeImg.File) })
 
 	// Recovery does not mount persistent, so try to mount it. Ignore errors, as it's not mandatory.
+	// This was used by luet extraction IIRC to not exhaust the /tmp dir
+	// Not sure if its on use anymore and we should drop it
+	// TODO: Check if we really need persistent mounted here
 	persistentPart := u.spec.Partitions.Persistent
 	if persistentPart != nil {
 		// Create the dir otherwise the check for mounted dir fails
@@ -162,8 +165,6 @@ func (u *UpgradeAction) Run() (err error) {
 			umount, err = e.MountRWPartition(persistentPart)
 			if err != nil {
 				u.config.Logger.Warn("could not mount persistent partition: %s", err.Error())
-			} else {
-				cleanup.Push(umount)
 			}
 		}
 	}
diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index cff10f9..a81d07d 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -125,9 +125,11 @@ func NewInstallElementalPartitions(spec *v1.InstallSpec) v1.ElementalPartitions
 		Flags:           []string{},
 	}
 
+	// Double the space for recovery, as upgrades use the recovery partition to create the transition image for upgrades
+	// so we need twice the space to do a proper upgrade
 	pt.Recovery = &v1.Partition{
 		FilesystemLabel: constants.RecoveryLabel,
-		Size:            spec.Recovery.Size + 200,
+		Size:            (spec.Recovery.Size * 2) + 200,
 		Name:            constants.RecoveryPartName,
 		FS:              constants.LinuxFs,
 		MountPoint:      constants.RecoveryDir,
@@ -136,9 +138,11 @@ func NewInstallElementalPartitions(spec *v1.InstallSpec) v1.ElementalPartitions
 
 	// Add 1 Gb to the partition so images can grow a bit, otherwise you are stuck with the smallest space possible and
 	// there is no coming back from that
+	// Also multiply the space for active, as upgrades use the state partition to create the transition image for upgrades
+	// so we need twice the space to do a proper upgrade
 	pt.State = &v1.Partition{
 		FilesystemLabel: constants.StateLabel,
-		Size:            spec.Active.Size + spec.Passive.Size + 1000,
+		Size:            (spec.Active.Size * 2) + spec.Passive.Size + 1000,
 		Name:            constants.StatePartName,
 		FS:              constants.LinuxFs,
 		MountPoint:      constants.StateDir,