Merge pull request #2442 from OSInside/disk_to_context_manager

Move Disk to context manager
OSInside · Jan 31, 2024 · 01acc87 · 01acc87
2 parents 82d9a6c + c5b3c83
commit 01acc87
Show file tree

Hide file tree

Showing 6 changed files with 256 additions and 184 deletions.
diff --git a/kiwi/bootloader/config/systemd_boot.py b/kiwi/bootloader/config/systemd_boot.py
@@ -18,6 +18,7 @@
 import os
 import glob
 from string import Template
+from contextlib import ExitStack
 from typing import Dict
 
 # project
@@ -119,46 +120,44 @@ def _create_embedded_fat_efi_image(self, path: str):
         )
         with LoopDevice(path) as loop_provider:
             loop_provider.create(overwrite=False)
-            disk = Disk('gpt', loop_provider)
-            disk.map_partitions()
-            disk.partitioner.partition_id = 1
-            disk._add_to_map('efi')
-            Command.run(
-                ['mkdosfs', '-n', 'BOOT', disk.partition_map['efi']]
-            )
-            Path.create(f'{self.root_dir}/boot/efi')
-            efi_mount = MountManager(
-                device=disk.partition_map['efi'],
-                mountpoint=f'{self.root_dir}/boot/efi'
-            )
-            device_mount = MountManager(
-                device='/dev',
-                mountpoint=f'{self.root_dir}/dev'
-            )
-            proc_mount = MountManager(
-                device='/proc',
-                mountpoint=f'{self.root_dir}/proc'
-            )
-            sys_mount = MountManager(
-                device='/sys',
-                mountpoint=f'{self.root_dir}/sys'
-            )
-            efi_mount.mount()
-            device_mount.bind_mount()
-            proc_mount.bind_mount()
-            sys_mount.bind_mount()
-            try:
-                self._run_bootctl(self.root_dir)
-                self.set_loader_entry(self.root_dir, self.target.live)
-            finally:
-                efi_mount.umount()
-                device_mount.umount()
-                proc_mount.umount()
-                sys_mount.umount()
-            Command.run(
-                ['dd', f'if={disk.partition_map["efi"]}', f'of={path}.img']
-            )
-            del disk
+            with Disk('gpt', loop_provider) as disk:
+                disk.map_partitions()
+                disk.partitioner.partition_id = 1
+                disk._add_to_map('efi')
+                Command.run(
+                    ['mkdosfs', '-n', 'BOOT', disk.partition_map['efi']]
+                )
+                Path.create(f'{self.root_dir}/boot/efi')
+                with ExitStack() as stack:
+                    efi_mount = MountManager(
+                        device=disk.partition_map['efi'],
+                        mountpoint=f'{self.root_dir}/boot/efi'
+                    )
+                    stack.push(efi_mount)
+                    device_mount = MountManager(
+                        device='/dev',
+                        mountpoint=f'{self.root_dir}/dev'
+                    )
+                    stack.push(device_mount)
+                    proc_mount = MountManager(
+                        device='/proc',
+                        mountpoint=f'{self.root_dir}/proc'
+                    )
+                    stack.push(proc_mount)
+                    sys_mount = MountManager(
+                        device='/sys',
+                        mountpoint=f'{self.root_dir}/sys'
+                    )
+                    stack.push(sys_mount)
+                    efi_mount.mount()
+                    device_mount.bind_mount()
+                    proc_mount.bind_mount()
+                    sys_mount.bind_mount()
+                    self._run_bootctl(self.root_dir)
+                    self.set_loader_entry(self.root_dir, self.target.live)
+                Command.run(
+                    ['dd', f'if={disk.partition_map["efi"]}', f'of={path}.img']
+                )
 
         Command.run(
             ['mv', f'{path}.img', path]

diff --git a/kiwi/builder/disk.py b/kiwi/builder/disk.py
@@ -336,40 +336,39 @@ def create_disk(self) -> Result:
             loop_provider.create()
 
             # create the disk partitioner, still unmapped
-            disk = self._disk_instance(loop_provider)
-
-            # create disk partitions and instance device map
-            device_map = self._build_and_map_disk_partitions(
-                disk, disksize_mbytes
-            )
-
-            # update device and disk id map in case of no root write partition
-            if self.root_filesystem_is_overlay and \
-               self.root_filesystem_has_write_partition is False:
-                device_map['root'] = device_map['readonly']
-                disk.public_partition_id_map['kiwi_RootPart'] = \
-                    disk.public_partition_id_map['kiwi_ROPart']
-
-            # create raid on current root device if requested
-            if self.mdraid:
-                raid_root = self._raid_instance(device_map)
-                device_map = self._map_raid(device_map, disk, raid_root)
-
-            # create integrity on current root device if requested
-            if self.integrity:
-                integrity_root = self._integrity_instance(device_map)
-                device_map = self._map_integrity(device_map, integrity_root)
-
-            # create luks on current root device if requested
-            if self.luks is not None:
-                luks_root = self._luks_instance(device_map)
-                device_map = self._map_luks(device_map, luks_root)
-
-            # create system layout for root system
-            self._create_system_instance(device_map)
+            with self._disk_instance(loop_provider) as disk:
+                # create disk partitions and instance device map
+                device_map = self._build_and_map_disk_partitions(
+                    disk, disksize_mbytes
+                )
 
-            # sync system data, configure system, setup loader and initrd
-            self._process_build(device_map, disk)
+                # update device and disk id map if no root write partition
+                if self.root_filesystem_is_overlay and \
+                   self.root_filesystem_has_write_partition is False:
+                    device_map['root'] = device_map['readonly']
+                    disk.public_partition_id_map['kiwi_RootPart'] = \
+                        disk.public_partition_id_map['kiwi_ROPart']
+
+                # create raid on current root device if requested
+                if self.mdraid:
+                    raid_root = self._raid_instance(device_map)
+                    device_map = self._map_raid(device_map, disk, raid_root)
+
+                # create integrity on current root device if requested
+                if self.integrity:
+                    integrity_root = self._integrity_instance(device_map)
+                    device_map = self._map_integrity(device_map, integrity_root)
+
+                # create luks on current root device if requested
+                if self.luks is not None:
+                    luks_root = self._luks_instance(device_map)
+                    device_map = self._map_luks(device_map, luks_root)
+
+                # create system layout for root system
+                self._create_system_instance(device_map)
+
+                # sync system data, configure system, setup loader and initrd
+                self._process_build(device_map, disk)
 
         # store image bundle_format in result
         if self.bundle_format:

diff --git a/kiwi/storage/disk.py b/kiwi/storage/disk.py
@@ -99,6 +99,9 @@ def __init__(
 
         self.table_type = table_type
 
+    def __enter__(self):
+        return self
+
     def get_device(self) -> Dict[str, MappedDevice]:
         """
         Names of partition devices
@@ -510,7 +513,7 @@ def _add_to_map(self, name):
             self.partition_map[name] = device_node
             self.partition_id_map[name] = partition_number
 
-    def __del__(self):
+    def __exit__(self, exc_type, exc_value, traceback):
         if self.storage_provider.is_loop() and self.is_mapped:
             log.info('Cleaning up %s instance', type(self).__name__)
             try:
@@ -524,8 +527,9 @@ def __del__(self):
                     Command.run(
                         ['partx', '--delete', self.storage_provider.get_device()]
                     )
-            except Exception:
-                log.warning(
-                    'cleanup of partition device maps failed, %s still busy',
-                    self.storage_provider.get_device()
+            except Exception as issue:
+                log.error(
+                    'cleanup of partition maps on {} failed with: {}'.format(
+                        self.storage_provider.get_device(), issue
+                    )
                 )
diff --git a/test/unit/bootloader/config/systemd_boot_test.py b/test/unit/bootloader/config/systemd_boot_test.py
@@ -194,7 +194,7 @@ def test_create_embedded_fat_efi_image(
     ):
         target = Mock()
         self.bootloader.target = target
-        mock_Disk.return_value.partition_map = {
+        mock_Disk.return_value.__enter__.return_value.partition_map = {
             'efi': 'efi_device'
         }
         self.bootloader._create_embedded_fat_efi_image('ESP')