*Always* use link_dm_node() when install grub 1.99

Grub install fails when in a chrooted environment,
it has nothing to do with the volume being a loopback volume

bootstrapvz/common/fs/__init__.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 
 
 def get_partitions():
@@ -16,7 +17,8 @@ def get_partitions():
 	return matches
 
 
-def remount(volume, fn):
+@contextmanager
+def unmounted(volume):
 	from bootstrapvz.base.fs.partitionmaps.none import NoPartitions
 
 	p_map = volume.partition_map
@@ -24,9 +26,8 @@ def remount(volume, fn):
 	p_map.root.unmount()
 	if not isinstance(p_map, NoPartitions):
 		p_map.unmap(volume)
-		result = fn()
+		yield
 		p_map.map(volume)
 	else:
-		result = fn()
+		yield
 	p_map.root.mount(destination=root_dir)
-	return result

bootstrapvz/common/tasks/grub.py

@@ -39,26 +39,16 @@ class InstallGrub_1_99(Task):
 
 	@classmethod
 	def run(cls, info):
-
-		from ..fs import remount
 		p_map = info.volume.partition_map
 
-		def link_fn():
+		# GRUB screws up when installing in chrooted environments
+		# so we fake a real harddisk with dmsetup.
+		# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
+		from ..fs import unmounted
+		with unmounted(info.volume):
 			info.volume.link_dm_node()
 			if isinstance(p_map, partitionmaps.none.NoPartitions):
 				p_map.root.device_path = info.volume.device_path
-
-		def unlink_fn():
-			info.volume.unlink_dm_node()
-			if isinstance(p_map, partitionmaps.none.NoPartitions):
-				p_map.root.device_path = info.volume.device_path
-
-		# GRUB cannot deal with installing to loopback devices
-		# so we fake a real harddisk with dmsetup.
-		# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
-		from ..fs.loopbackvolume import LoopbackVolume
-		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume, link_fn)
 		try:
 			[device_path] = log_check_call(['readlink', '-f', info.volume.device_path])
 			device_map_path = os.path.join(info.root, 'boot/grub/device.map')
@@ -77,13 +67,11 @@ class InstallGrub_1_99(Task):
 			# Install grub
 			log_check_call(['chroot', info.root, 'grub-install', device_path])
 			log_check_call(['chroot', info.root, 'update-grub'])
-		except Exception:
-			if isinstance(info.volume, LoopbackVolume):
-				remount(info.volume, unlink_fn)
-			raise
-
-		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume, unlink_fn)
+		finally:
+			with unmounted(info.volume):
+				info.volume.unlink_dm_node()
+				if isinstance(p_map, partitionmaps.none.NoPartitions):
+					p_map.root.device_path = info.volume.device_path
 
 
 class InstallGrub_2(Task):

bootstrapvz/plugins/prebootstrapped/tasks.py

@@ -4,7 +4,7 @@ from bootstrapvz.common.tasks import volume
 from bootstrapvz.common.tasks import packages
 from bootstrapvz.providers.virtualbox.tasks import guest_additions
 from bootstrapvz.providers.ec2.tasks import ebs
-from bootstrapvz.common.fs import remount
+from bootstrapvz.common.fs import unmounted
 from shutil import copyfile
 import os.path
 import time
@@ -19,9 +19,9 @@ class Snapshot(Task):
 
 	@classmethod
 	def run(cls, info):
-		def mk_snapshot():
-			return info.volume.snapshot()
-		snapshot = remount(info.volume, mk_snapshot)
+		snapshot = None
+		with unmounted(info.volume):
+			snapshot = info.volume.snapshot()
 		msg = 'A snapshot of the bootstrapped volume was created. ID: ' + snapshot.id
 		log.info(msg)
 
@@ -55,9 +55,8 @@ class CopyImage(Task):
 		loopback_backup_name = 'volume-{id}.{ext}.backup'.format(id=info.run_id, ext=info.volume.extension)
 		destination = os.path.join(info.manifest.bootstrapper['workspace'], loopback_backup_name)
 
-		def mk_snapshot():
+		with unmounted(info.volume):
 			copyfile(info.volume.image_path, destination)
-		remount(info.volume, mk_snapshot)
 		msg = 'A copy of the bootstrapped volume was created. Path: ' + destination
 		log.info(msg)