Fix dual-booting (#1250)

# Fixes

 * Optimized partition lookups
 * Fixed re-use of partition UUID's
 * `BlockDevice().get_partition()` now supports looking up both `PARTUUID` and `UUID` for a partition under itself
 * Partitions listed in `--disk-layout` that doesn't have a PARTUUID/UUID should no longer cause an exception, but instead logs a warning and they will simply be ignored
 * `Filesystem().add_partition()` now handles `DiskError` raised by `partition.part_uuid`
 * Fixed issue on normal partitions where the device was not properly frozen in `lambda` calls, meaning two or more mount-points shared the same `device_instance`.
 * Lowered global `DISK_RETRY_ATTEMPTS` to 5, as the timeouts are linear *(`range(DISK_RETRY_ATTEMPTS) * DISK_TIMEOUTS`)*

3 files changed, 55 insertions, 35 deletions

diff --git a/archinstall/lib/disk/blockdevice.py b/archinstall/lib/disk/blockdevice.py
index 15f03789..c4707abd 100644
--- a/archinstall/lib/disk/blockdevice.py
+++ b/archinstall/lib/disk/blockdevice.py
@@ -289,19 +289,27 @@ class BlockDevice:
 	def flush_cache(self) -> None:
 		self.part_cache = {}
 
-	def get_partition(self, uuid :str) -> Partition:
-		count = 0
-		while count < 5:
-			for partition_uuid, partition in self.partitions.items():
-				if partition.part_uuid.lower() == uuid.lower():
-					return partition
-			else:
-				log(f"uuid {uuid} not found. Waiting for {count +1} time",level=logging.DEBUG)
-				time.sleep(float(storage['arguments'].get('disk-sleep', 0.2)))
-				count += 1
-		else:
-			log(f"Could not find {uuid} in disk after 5 retries",level=logging.INFO)
-			print(f"Cache: {self.part_cache}")
-			print(f"Partitions: {self.partitions.items()}")
-			print(f"UUID: {[uuid]}")
-			raise DiskError(f"New partition {uuid} never showed up after adding new partition on {self}")
+	def get_partition(self, uuid :Optional[str] = None, partuuid :Optional[str] = None) -> Partition:
+		if not uuid and not partuuid:
+			raise ValueError(f"BlockDevice.get_partition() requires either a UUID or a PARTUUID for lookups.")
+
+		for count in range(storage.get('DISK_RETRY_ATTEMPTS', 5)):
+			for partition_index, partition in self.partitions.items():
+				try:
+					if uuid and partition.uuid.lower() == uuid.lower():
+						return partition
+					elif partuuid and partition.part_uuid.lower() == partuuid.lower():
+						return partition
+				except DiskError as error:
+					# Most likely a blockdevice that doesn't support or use UUID's
+					# (like Microsoft recovery partition)
+					log(f"Could not get UUID/PARTUUID of {partition}: {error}", level=logging.DEBUG, fg="gray")
+					pass
+
+			log(f"uuid {uuid} or {partuuid} not found. Waiting {storage.get('DISK_TIMEOUTS', 1) * count}s for next attempt",level=logging.DEBUG)
+			time.sleep(storage.get('DISK_TIMEOUTS', 1) * count)
+				
+		log(f"Could not find {uuid}/{partuuid} in disk after 5 retries", level=logging.INFO)
+		log(f"Cache: {self.part_cache}")
+		log(f"Partitions: {self.partitions.items()}")
+		raise DiskError(f"New partition {uuid}/{partuuid} never showed up after adding new partition on {self}")
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index 8d8f596e..66112ccc 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -97,7 +97,9 @@ class Filesystem:
 				print(_("Re-using partition instance: {}").format(partition_instance))
 				partition['device_instance'] = partition_instance
 			else:
-				raise ValueError(f"{self}.load_layout() doesn't know how to continue without a new partition definition or a UUID ({partition.get('PARTUUID')}) on the device ({self.blockdevice.get_partition(uuid=partition.get('PARTUUID'))}).")
+				log(f"{self}.load_layout() doesn't know how to work without 'wipe' being set or UUID ({partition.get('PARTUUID')}) was given and found.", fg="yellow", level=logging.WARNING)
+				continue
+				# raise ValueError(f"{self}.load_layout() doesn't know how to continue without a new partition definition or a UUID ({partition.get('PARTUUID')}) on the device ({self.blockdevice.get_partition(uuid=partition.get('PARTUUID'))}).")
 
 			if partition.get('filesystem', {}).get('format', False):
 
@@ -206,7 +208,12 @@ class Filesystem:
 	def add_partition(self, partition_type :str, start :str, end :str, partition_format :Optional[str] = None) -> Partition:
 		log(f'Adding partition to {self.blockdevice}, {start}->{end}', level=logging.INFO)
 
-		previous_partition_uuids = {partition.part_uuid for partition in self.blockdevice.partitions.values()}
+		previous_partuuids = []
+		for partition in self.blockdevice.partitions.values():
+			try:
+				previous_partuuids.append(partition.part_uuid)
+			except DiskError:
+				pass
 
 		if self.mode == MBR:
 			if len(self.blockdevice.partitions) > 3:
@@ -220,36 +227,41 @@ class Filesystem:
 		log(f"Adding partition using the following parted command: {parted_string}", level=logging.DEBUG)
 
 		if self.parted(parted_string):
-			count = 0
-			while count < 10:
-				new_uuid = None
-				new_uuid_set = (previous_partition_uuids ^ {partition.part_uuid for partition in self.blockdevice.partitions.values()})
+			for count in range(storage.get('DISK_RETRY_ATTEMPTS', 3)):
+				self.partprobe()
 
-				if len(new_uuid_set) > 0:
-					new_uuid = new_uuid_set.pop()
+				new_partition_uuids = []
+				for partition in self.blockdevice.partitions.values():
+					try:
+						new_partition_uuids.append(partition.part_uuid)
+					except DiskError:
+						pass
+
+				new_partuuid_set = (set(previous_partuuids) ^ set(new_partition_uuids))
+
+				print(previous_partuuids, new_partition_uuids, new_partuuid_set)
 
-				if new_uuid:
+				if len(new_partuuid_set) and (new_partuuid := new_partuuid_set.pop()):
 					try:
-						return self.blockdevice.get_partition(new_uuid)
+						return self.blockdevice.get_partition(partuuid=new_partuuid)
 					except Exception as err:
 						log(f'Blockdevice: {self.blockdevice}', level=logging.ERROR, fg="red")
 						log(f'Partitions: {self.blockdevice.partitions}', level=logging.ERROR, fg="red")
-						log(f'Partition set: {new_uuid_set}', level=logging.ERROR, fg="red")
-						log(f'New UUID: {[new_uuid]}', level=logging.ERROR, fg="red")
+						log(f'Partition set: {new_partuuid_set}', level=logging.ERROR, fg="red")
+						log(f'New UUID: {[new_partuuid]}', level=logging.ERROR, fg="red")
 						log(f'get_partition(): {self.blockdevice.get_partition}', level=logging.ERROR, fg="red")
 						raise err
 				else:
-					count += 1
-					log(f"Could not get UUID for partition. Waiting before retry attempt {count} of 10 ...",level=logging.DEBUG)
-					time.sleep(float(storage['arguments'].get('disk-sleep', 0.2)))
+					log(f"Could not get UUID for partition. Waiting {storage.get('DISK_TIMEOUTS', 1) * count}s before retrying.",level=logging.DEBUG)
+					time.sleep(storage.get('DISK_TIMEOUTS', 1) * count)
 			else:
 				log("Add partition is exiting due to excessive wait time", level=logging.ERROR, fg="red")
 				raise DiskError(f"New partition never showed up after adding new partition on {self}.")
 
 		# TODO: This should never be able to happen
 		log(f"Could not find the new PARTUUID after adding the partition.", level=logging.ERROR, fg="red")
-		log(f"Previous partitions: {previous_partition_uuids}", level=logging.ERROR, fg="red")
-		log(f"New partitions: {(previous_partition_uuids ^ {partition.part_uuid for partition in self.blockdevice.partitions.values()})}", level=logging.ERROR, fg="red")
+		log(f"Previous partitions: {previous_partuuids}", level=logging.ERROR, fg="red")
+		log(f"New partitions: {(previous_partuuids ^ {partition.part_uuid for partition in self.blockdevice.partitions.values()})}", level=logging.ERROR, fg="red")
 		raise DiskError(f"Could not add partition using: {parted_string}")
 
 	def set_name(self, partition: int, name: str) -> bool:
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index e33c600c..bb1ffeb6 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -101,7 +101,7 @@ class Partition:
 
 		except SysCallError as error:
 			# Not mounted anywhere most likely
-			log(f"Could not locate mount information for {self.path}: {error}", level=logging.DEBUG)
+			log(f"Could not locate mount information for {self.path}: {error}", level=logging.DEBUG, fg="grey")
 			pass
 
 		return None
@@ -216,7 +216,7 @@ class Partition:
 			if not self.partprobe():
 				raise DiskError(f"Could not perform partprobe on {self.device_path}")
 
-			time.sleep(max(0.1, storage['DISK_TIMEOUTS'] * i))
+			time.sleep(storage.get('DISK_TIMEOUTS', 1) * i)
 
 			partuuid = self._safe_uuid
 			if partuuid: