storage: look harder for the gap after removing a partition

When removing a partition, we expect the resulting gap to be either
before (if there was free space) or exactly where the partition used to
be.

However, in practice, Subiquity sometimes places the gap slightly after
the partition. This is probably a bug but for now let's look harder for
the resulting gap.

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

subiquity/common/filesystem/gaps.py

@@ -14,6 +14,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import functools
+import logging
 from typing import List, Tuple
 
 import attr
@@ -30,6 +31,8 @@
     align_up,
 )
 
+log = logging.getLogger("subiquity.common.filesystem.gaps")
+
 
 # should also set on_setattr=None with attrs 20.1.0
 @attr.s(auto_attribs=True, frozen=True)
@@ -314,10 +317,11 @@ def at_offset(device, offset):
     return None
 
 
-def after(device, offset):
-    """Find the first gap that is after this offset."""
+def after(device, offset, *, only_gap=True):
+    """Find the first gap that is after this offset. If only_gap is False, find
+    the first gap (or partition!) that is after this offset."""
     for pg in parts_and_gaps(device):
-        if isinstance(pg, Gap):
+        if isinstance(pg, Gap) or not only_gap:
             if pg.offset > offset:
                 return pg
     return None
@@ -331,3 +335,43 @@ def includes(device, offset):
         if pg.offset <= offset < (pg.offset + pg.size):
             return pg
     return None
+
+
+def find_gap_after_removal(disk: Disk, removed_partition: Partition) -> Gap:
+    """Assuming that the specified partition was removed from the disk, return
+    the corresponding, resulting gap that "contains" the partition. The offset
+    of the resulting gap can either:
+    * match that of the removed partition
+    * be lower than that of the removed partition (we automatically claim free
+    space prepending the removed partition)
+    * be greater (surprisingly) than that of the removed partition. This is
+    probably a bug."""
+    gap = includes(disk, removed_partition.offset)
+
+    if gap is not None:
+        return gap
+
+    # Sometimes, after removal, the gap will appear slighly *after* where the
+    # partition used to be. This is notably an issue when dealing with logical
+    # partitions.
+    # This is arguably a bug (i.e., we probably block too much space) but let's
+    # try to look harder for the gap since we know it can happen.
+    part_or_gap = after(disk, removed_partition.offset, only_gap=False)
+
+    # Ensure that what we find is actually a gap and ensure that it covers at
+    # least one part of the partition we removed.
+    if (
+        part_or_gap is None
+        or not isinstance(part_or_gap, Gap)
+        or part_or_gap.offset >= removed_partition.offset + removed_partition.size
+    ):
+        raise RuntimeError(
+            f"could not find a gap after removing {removed_partition.id}"
+        )
+
+    gap = part_or_gap
+    log.debug(
+        "no gap found on %s that includes offset %s", disk.id, removed_partition.offset
+    )
+    log.debug("using gap at offset %s instead", gap.offset)
+    return gap

subiquity/common/filesystem/tests/test_gaps.py

@@ -267,6 +267,108 @@ def test_half_allocated(self):
         self.assertEqual(gap, gaps.includes(d, gap.offset + gap.size // 2))
 
 
+class TestFindGapAfterRemoval(unittest.TestCase):
+    def test_no_free_space_single_part(self):
+        m, d = make_model_and_disk(ptable="gpt")
+        p = make_partition(m, d, size=-1)
+
+        m.remove_partition(p)
+
+        gap = gaps.find_gap_after_removal(d, p)
+
+        self.assertEqual(p.offset, gap.offset)
+        self.assertEqual(p.size, gap.size)
+
+    def test_no_free_space_two_parts(self):
+        m, d = make_model_and_disk(ptable="gpt")
+        make_partition(m, d)
+        p2 = make_partition(m, d, size=-1)
+
+        m.remove_partition(p2)
+
+        gap = gaps.find_gap_after_removal(d, p2)
+
+        self.assertEqual(p2.offset, gap.offset)
+        self.assertEqual(p2.size, gap.size)
+
+    def test_free_space_before(self):
+        m, d = make_model_and_disk(ptable="gpt")
+        p = make_partition(m, d, offset=5 * MiB, size=10 * MiB)
+
+        m.remove_partition(p)
+
+        gap = gaps.find_gap_after_removal(d, p)
+
+        self.assertEqual(d.alignment_data().min_start_offset, gap.offset)
+        self.assertLess(gap.offset, p.offset)
+
+    def test_unaligned_part(self):
+        m, d = make_model_and_disk(ptable="gpt")
+        make_partition(m, d, offset=1 * MiB, size=int(10.5 * MiB))
+        p2 = make_partition(m, d, offset=int(11.5 * MiB), size=10 * MiB)
+
+        m.remove_partition(p2)
+
+        gap = gaps.find_gap_after_removal(d, p2)
+
+        self.assertGreater(gap.offset, p2.offset)
+
+    def test_unaligned_extended(self):
+        """This test reuses data from threebuntu-on-msdos.json ; where the
+        logical partitions are aligned to the MiB boundary, but the extended
+        partition is not. When removing the first logical partition, the
+        resulting gap does not appear where the partition used to be, but
+        after. This is arguably a bug and with this test-case, we are
+        acknowledging that this quirk exists.
+        disk
+        ----
+        { "sectorsize": 512, "unit": "sectors", "size": "161061273600" }
+        partitions
+        ----------
+         { "node": "/dev/sda1", "size": 97654784, "start": 2048, "type": "83" },
+         { "node": "/dev/sda2", "size": 215863298, "start": 97658878, "type": "5" },
+         { "node": "/dev/sda3", "size": 1048576, "start": 313522176, "type": "ef" },
+         { "node": "/dev/sda5", "size": 97654784, "start": 97658880, "type": "83" },
+         { "node": "/dev/sda6", "size": 118206464, "start": 195315712, "type": "83" },
+        """
+
+        def make_part(*, offset, size, sector_size=512, **kwargs):
+            return make_partition(
+                m,
+                d,
+                preserve=True,
+                size=size * sector_size,
+                offset=offset * sector_size,
+                **kwargs,
+            )
+
+        m, d = make_model_and_disk(size=161061273600, ptable="dos", storage_version=2)
+        p1 = make_part(offset=2048, size=97654784)
+        p2 = make_part(offset=97658878, size=215863298, flag="extended")
+        p3 = make_part(offset=313522176, size=1048576)
+        p5 = make_part(offset=97658880, size=97654784, flag="logical")
+        p6 = make_part(offset=195315712, size=118206464, flag="logical")
+
+        # Just so we can use MiB later without thinking.
+        self.assertEqual(MiB, d.alignment_data().part_align)
+
+        self.assertTrue(p1.offset % MiB == 0)
+        self.assertTrue(p3.offset % MiB == 0)
+        self.assertTrue(p5.offset % MiB == 0)
+        self.assertTrue(p6.offset % MiB == 0)
+
+        # The extended partition is unaligned
+        self.assertFalse(p2.offset % MiB == 0)
+
+        m.remove_partition(p5)
+
+        gap = gaps.find_gap_after_removal(d, p5)
+
+        # Ideally, the gap should be at p5.offset, not after p5.offset.
+        # Let's change this to assertEqual if we fix the "issue".
+        self.assertGreater(gap.offset, p5.offset)
+
+
 class TestDiskGaps(unittest.TestCase):
     def test_no_partition_gpt(self):
         size = 1 << 30

subiquity/models/tests/test_filesystem.py

@@ -167,8 +167,8 @@ def make_disk(fs_model=None, **kw):
     return disk
 
 
-def make_model_and_disk(bootloader=None, **kw):
-    model = make_model(bootloader)
+def make_model_and_disk(bootloader=None, storage_version=None, **kw):
+    model = make_model(bootloader, storage_version)
     return model, make_disk(model, **kw)
 
 

subiquity/server/controllers/filesystem.py

@@ -742,7 +742,7 @@ def start_guided_erase_install(
         will not necessarily match partition.offset and partition.size."""
         partition = self.get_partition(disk, target.partition_number)
         self.delete_partition(partition, override_preserve=True)
-        return gaps.includes(disk, partition.offset)
+        return gaps.find_gap_after_removal(disk, removed_partition=partition)
 
     def set_info_for_capability(self, capability: GuidedCapability):
         """Given a request for a capability, select the variation to use."""
@@ -1304,7 +1304,9 @@ def available_erase_install_scenarios(
                 if not boot.can_be_boot_device(altered_disk, with_reformatting=False):
                     continue
 
-                gap = gaps.includes(altered_disk, offset=partition.offset)
+                gap = gaps.find_gap_after_removal(
+                    altered_disk, removed_partition=partition
+                )
                 if not self.use_gap_has_enough_room_for_partitions(altered_disk, gap):
                     log.error(
                         "skipping TargetEraseInstall: not enough room for primary"