Way better rollback architecture through improvements in flexibility

base/main.py

@@ -1,3 +1,5 @@
+import logging
+log = logging.getLogger(__name__)
 
 
 def main():
@@ -29,4 +31,17 @@ def run(args):
 
 	from bootstrapinfo import BootstrapInformation
 	bootstrap_info = BootstrapInformation(manifest=manifest, debug=args.debug)
-	tasklist.run(bootstrap_info)
+
+	try:
+		tasklist.run(bootstrap_info)
+	except Exception as e:
+		log.exception(e)
+		log.error('Rolling back')
+		rollback_tasklist = TaskList()
+		provider.rollback_tasks(rollback_tasklist, tasklist.tasks_completed, manifest)
+		for plugin in manifest.loaded_plugins:
+			rollback_tasks = getattr(plugin, 'rollback_tasks', None)
+			if callable(rollback_tasks):
+				plugin.rollback_tasks(rollback_tasklist, tasklist.tasks_completed, manifest)
+		rollback_tasklist.run(bootstrap_info)
+		log.info('Successfully completed rollback')

base/tasklist.py

@@ -7,6 +7,7 @@ class TaskList(object):
 
 	def __init__(self):
 		self.tasks = set()
+		self.tasks_completed = []
 
 	def add(self, *args):
 		self.tasks.update(args)
@@ -25,28 +26,13 @@ class TaskList(object):
 		task_list = self.create_list(self.tasks)
 		log.debug('Tasklist:\n\t{list}'.format(list='\n\t'.join(repr(task) for task in task_list)))
 
-		tasks_completed = []
-		try:
-			for task in task_list:
-				if hasattr(task, 'description'):
-					log.info(task.description)
-				else:
-					log.info('Running {task}'.format(task=task))
-				task.run(bootstrap_info)
-				tasks_completed.append(task)
-		except Exception as e:
-			log.exception(e)
-			log.error('Rolling back')
-			for task in reversed(tasks_completed):
-				rollback = getattr(task, 'rollback', None)
-				if not callable(rollback):
-					continue
-				if hasattr(task, 'rollback_description'):
-					log.warning(task.rollback_description)
-				else:
-					log.warning('Rolling back {task}'.format(task=task))
-				task.rollback(bootstrap_info)
-			log.info('Successfully completed rollback')
+		for task in task_list:
+			if hasattr(task, 'description'):
+				log.info(task.description)
+			else:
+				log.info('Running {task}'.format(task=task))
+			task.run(bootstrap_info)
+			self.tasks_completed.append(task)
 
 	def create_list(self, tasks):
 		from common.phases import order

common/phases.py

@@ -4,21 +4,21 @@ preparation = Phase('Initializing connections, fetching data etc.')
 volume_creation = Phase('Creating the volume to bootstrap onto')
 volume_preparation = Phase('Formatting the bootstrap volume')
 volume_mounting = Phase('Mounting bootstrap volume')
-install_os = Phase('Installing the operating system')
-modify_system = Phase('Installing software, modifying configuration files etc.')
-clean_system = Phase('Removing sensitive data, temporary files and other leftovers')
-unmount_volume = Phase('Unmounting the bootstrap volume')
-register_image = Phase('Uploading/Registering with the provider')
-cleanup = Phase('Removing temporary files')
+os_installation = Phase('Installing the operating system')
+system_modification = Phase('Installing software, modifying configuration files etc.')
+system_cleaning = Phase('Removing sensitive data, temporary files and other leftovers')
+volume_unmounting = Phase('Unmounting the bootstrap volume')
+image_registration = Phase('Uploading/Registering with the provider')
+cleaning = Phase('Removing temporary files')
 
 order = [preparation,
          volume_creation,
          volume_preparation,
          volume_mounting,
-         install_os,
-         modify_system,
-         clean_system,
-         unmount_volume,
-         register_image,
-         cleanup,
+         os_installation,
+         system_modification,
+         system_cleaning,
+         volume_unmounting,
+         image_registration,
+         cleaning,
          ]

common/tasks.py

@@ -3,7 +3,7 @@ import phases
 
 
 class TriggerRollback(Task):
-	phase = phases.cleanup
+	phase = phases.cleaning
 
 	description = 'Triggering a rollback by throwing an exception'
 

providers/ec2/__init__.py

@@ -1,11 +1,11 @@
 from manifest import Manifest
+from tasks import packages
+from tasks import connection
+from tasks import host
+from tasks import ebs
 
 
 def tasks(tasklist, manifest):
-	from tasks import packages
-	from tasks import connection
-	from tasks import host
-	from tasks import ebs
 	tasklist.add(packages.HostPackages(), packages.ImagePackages(), host.CheckPackages(),
 	             connection.GetCredentials(), host.GetInfo(), connection.Connect())
 	if manifest.volume['backing'].lower() == 'ebs':
@@ -13,3 +13,12 @@ def tasks(tasklist, manifest):
 
 	from common.tasks import TriggerRollback
 	tasklist.add(TriggerRollback())
+
+
+def rollback_tasks(tasklist, tasks_completed, manifest):
+	completed = [type(task) for task in tasks_completed]
+	if manifest.volume['backing'].lower() == 'ebs':
+		if ebs.CreateVolume in completed and ebs.DeleteVolume not in completed:
+			tasklist.add(ebs.DeleteVolume())
+		if ebs.AttachVolume in completed and ebs.DetachVolume not in completed:
+			tasklist.add(ebs.DetachVolume())

providers/ec2/tasks/ebs.py

@@ -54,14 +54,30 @@ class AttachVolume(Task):
 			time.sleep(2)
 			info.volume.update()
 
-	rollback_description = 'Detaching the EBS volume'
 
-	def rollback(self, info):
+class DetachVolume(Task):
+	phase = phases.volume_unmounting
+	after = []
+
+	description = 'Detaching the EBS volume'
+
+	def run(self, info):
 		info.volume.detach()
 		while info.volume.attachment_state() is not None:
 			time.sleep(2)
 			info.volume.update()
 
 
+class DeleteVolume(Task):
+	phase = phases.cleaning
+	after = []
+
+	description = 'Deleting the EBS volume'
+
+	def run(self, info):
+		info.volume.delete()
+		del info.volume
+
+
 class VolumeError(TaskException):
 	pass