storage: don't let probing discard storage config in the middle of partitioning

subiquity/common/apidef.py

@@ -277,6 +277,10 @@ def GET(wait: bool = False, use_cached_result: bool = False) \
 
         def POST(config: Payload[list]): ...
 
+        class dry_run_wait_probe:
+            """This endpoint only works in dry-run mode."""
+            def POST() -> None: ...
+
         class reset:
             def POST() -> StorageResponse: ...
 
@@ -302,6 +306,16 @@ def POST(data: Payload[GuidedChoiceV2]) \
             class reset:
                 def POST() -> StorageResponseV2: ...
 
+            class ensure_transaction:
+                """This call will ensure that a transaction is initiated.
+                During a transaction, storage probing runs are not permitted to
+                reset the partitioning configuration.
+                A transaction will also be initiated by any v2_storage POST
+                request that modifies the partitioning configuration (e.g.,
+                add_partition, edit_partition, ...) but ensure_transaction can
+                be called early if desired. """
+                def POST() -> None: ...
+
             class reformat_disk:
                 def POST(data: Payload[ReformatDisk]) \
                     -> StorageResponseV2: ...

subiquity/server/controllers/filesystem.py

@@ -22,7 +22,7 @@
 import pathlib
 import select
 import time
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from curtin.commands.extract import AbstractSourceHandler
 from curtin.storage_config import ptable_uuid_to_flag_entry
@@ -160,6 +160,10 @@ def __init__(self, app):
         self._role_to_device: Dict[str: _Device] = {}
         self._device_to_structure: Dict[_Device: snapdapi.OnVolume] = {}
         self.use_tpm: bool = False
+        self.locked_probe_data = False
+        # If probe data come in while we are doing partitioning, store it in
+        # this variable. It will be picked up on next reset will pick it up.
+        self.queued_probe_data: Optional[Dict[str, Any]] = None
 
     def is_core_boot_classic(self):
         return self._system is not None
@@ -680,9 +684,22 @@ async def v2_orig_config_GET(self) -> StorageResponseV2:
 
     async def v2_reset_POST(self) -> StorageResponseV2:
         log.info("Resetting Filesystem model")
-        self.model.reset()
+        # From the API standpoint, it seems sound to set locked_probe_data back
+        # to False after a reset. But in practise, v2_reset_POST can be called
+        # during manual partitioning ; and we don't want to reenable automatic
+        # loading of probe data. Going forward, this could be controlled by an
+        # optional parameter maybe?
+        if self.queued_probe_data is not None:
+            log.debug("using newly obtained probe data")
+            self.model.load_probe_data(self.queued_probe_data)
+            self.queued_probe_data = None
+        else:
+            self.model.reset()
         return await self.v2_GET()
 
+    async def v2_ensure_transaction_POST(self) -> None:
+        self.locked_probe_data = True
+
     def get_capabilities(self):
         if self.is_core_boot_classic():
             classic_capabilities = []
@@ -772,16 +789,19 @@ async def v2_guided_GET(self, wait: bool = False) \
     async def v2_guided_POST(self, data: GuidedChoiceV2) \
             -> GuidedStorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         await self.guided(data)
         return await self.v2_guided_GET()
 
     async def v2_reformat_disk_POST(self, data: ReformatDisk) \
             -> StorageResponseV2:
+        self.locked_probe_data = True
         self.reformat(self.model._one(id=data.disk_id), data.ptable)
         return await self.v2_GET()
 
     async def v2_add_boot_partition_POST(self, disk_id: str) \
             -> StorageResponseV2:
+        self.locked_probe_data = True
         disk = self.model._one(id=disk_id)
         if boot.is_boot_device(disk):
             raise ValueError('device already has bootloader partition')
@@ -793,6 +813,7 @@ async def v2_add_boot_partition_POST(self, disk_id: str) \
     async def v2_add_partition_POST(self, data: AddPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         if data.partition.boot is not None:
             raise ValueError('add_partition does not support changing boot')
         disk = self.model._one(id=data.disk_id)
@@ -814,6 +835,7 @@ async def v2_add_partition_POST(self, data: AddPartitionV2) \
     async def v2_delete_partition_POST(self, data: ModifyPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         disk = self.model._one(id=data.disk_id)
         partition = self.get_partition(disk, data.partition.number)
         self.delete_partition(partition)
@@ -822,6 +844,7 @@ async def v2_delete_partition_POST(self, data: ModifyPartitionV2) \
     async def v2_edit_partition_POST(self, data: ModifyPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         disk = self.model._one(id=data.disk_id)
         partition = self.get_partition(disk, data.partition.number)
         if data.partition.size not in (None, partition.size) \
@@ -843,6 +866,15 @@ async def v2_edit_partition_POST(self, data: ModifyPartitionV2) \
         self.partition_disk_handler(disk, spec, partition=partition)
         return await self.v2_GET()
 
+    async def dry_run_wait_probe_POST(self) -> None:
+        if not self.app.opts.dry_run:
+            raise NotImplementedError
+
+        # This will start the probe task if not yet started.
+        self.ensure_probing()
+
+        await self._probe_task.task
+
     @with_context(name='probe_once', description='restricted={restricted}')
     async def _probe_once(self, *, context, restricted):
         if restricted:
@@ -867,7 +899,11 @@ async def _probe_once(self, *, context, restricted):
         with open(fpath, 'w') as fp:
             json.dump(storage, fp, indent=4)
         self.app.note_file_for_apport(key, fpath)
-        self.model.load_probe_data(storage)
+        if not self.locked_probe_data:
+            self.queued_probe_data = None
+            self.model.load_probe_data(storage)
+        else:
+            self.queued_probe_data = storage
 
     @with_context()
     async def _probe(self, *, context=None):
@@ -1022,6 +1058,14 @@ def stop_listening_udev(self):
         loop = asyncio.get_running_loop()
         loop.remove_reader(self._monitor.fileno())
 
+    def ensure_probing(self):
+        try:
+            self._probe_task.start_sync()
+        except TaskAlreadyRunningError:
+            log.debug('Skipping run of Probert - probe run already active')
+        else:
+            log.debug('Triggered Probert run on udev event')
+
     def _udev_event(self):
         cp = run_command(['udevadm', 'settle', '-t', '0'])
         if cp.returncode != 0:
@@ -1038,12 +1082,7 @@ def _udev_event(self):
         while select.select([self._monitor.fileno()], [], [], 0)[0]:
             action, dev = self._monitor.receive_device()
             log.debug("_udev_event %s %s", action, dev)
-        try:
-            self._probe_task.start_sync()
-        except TaskAlreadyRunningError:
-            log.debug('Skipping run of Probert - probe run already active')
-        else:
-            log.debug('Triggered Probert run on udev event')
+        self.ensure_probing()
 
     def make_autoinstall(self):
         rendered = self.model.render()