Merge pull request #207 from jwendell/deregister_ami

ec2_launch: Allow to deregister the AMI after launching image

bootstrapvz/plugins/ec2_launch/__init__.py

@@ -9,3 +9,5 @@ def resolve_tasks(taskset, manifest):
     taskset.add(tasks.LaunchEC2Instance)
     if 'print_public_ip' in manifest.plugins['ec2_launch']:
         taskset.add(tasks.PrintPublicIPAddress)
+    if manifest.plugins['ec2_launch'].get('deregister_ami', False):
+        taskset.add(tasks.DeregisterAMI)

bootstrapvz/plugins/ec2_launch/manifest-schema.yml

@@ -16,3 +16,5 @@ properties:
           instance_type: {type: string}
           print_public_ip: {type: string}
           tags: {type: object}
+          deregister_ami: {type: boolean}
+        additionalProperties: false

bootstrapvz/plugins/ec2_launch/tasks.py

@@ -1,10 +1,20 @@
 from bootstrapvz.base import Task
 from bootstrapvz.common import phases
 from bootstrapvz.providers.ec2.tasks import ami
-import time
 import logging
 
 
+# TODO: Merge with the method available in wip-integration-tests branch
+def waituntil(predicate, timeout=5, interval=0.05):
+	import time
+	threshhold = time.time() + timeout
+	while time.time() < threshhold:
+		if predicate():
+			return True
+		time.sleep(interval)
+	return False
+
+
 class LaunchEC2Instance(Task):
     description = 'Launching EC2 instance'
     phase = phases.image_registration
@@ -16,14 +26,14 @@ class LaunchEC2Instance(Task):
         r = conn.run_instances(info._ec2['image'],
                                security_group_ids=info.manifest.plugins['ec2_launch'].get('security_group_ids'),
                                instance_type=info.manifest.plugins['ec2_launch'].get('instance_type', 't2.micro'))
-        info._ec2['instance_id'] = r.instances[0].id
+        info._ec2['instance'] = r.instances[0]
 
         if 'tags' in info.manifest.plugins['ec2_launch']:
             def apply_format(v):
                 return v.format(**info.manifest_vars)
             tags = info.manifest.plugins['ec2_launch']['tags']
             r = {k: apply_format(v) for k, v in tags.items()}
-            conn.create_tags([info._ec2['instance_id']], r)
+            conn.create_tags([info._ec2['instance'].id], r)
 
 
 class PrintPublicIPAddress(Task):
@@ -40,31 +50,36 @@ class PrintPublicIPAddress(Task):
             filename = '/dev/null'
         f = open(filename, 'w')
 
-        i = 0
+        def instance_has_ip():
-        instance = None
+            ec2['instance'].update()
-        while True:
+            return ec2['instance'].ip_address
-            logger.debug('Waiting a bit to get instance metadata...')
-            time.sleep(5)
 
-            i += 1
+        if waituntil(instance_has_ip, timeout=120, interval=5):
-            if i > 10:
+            logger.info('******* EC2 IP ADDRESS: %s *******' % ec2['instance'].ip_address)
-                logger.error('Waited too much, giving up')
+            f.write(ec2['instance'].ip_address)
-                break
-
-            r = ec2['connection'].get_only_instances([ec2['instance_id']])
-            if not r and not r[0]:
-                logger.error('Could not get instance metadata')
-                break
-
-            instance = r[0]
-            if instance.ip_address:
-                break
-
-        if instance and instance.ip_address:
-            logger.info('******* EC2 IP ADDRESS: %s *******' % instance.ip_address)
-            f.write(instance.ip_address)
         else:
             logger.error('Could not get IP address for the instance')
             f.write('')
 
         f.close()
+
+
+class DeregisterAMI(Task):
+    description = 'Deregistering AMI'
+    phase = phases.image_registration
+    predecessors = [LaunchEC2Instance]
+
+    @classmethod
+    def run(cls, info):
+        ec2 = info._ec2
+        logger = logging.getLogger(__name__)
+
+        def instance_running():
+            ec2['instance'].update()
+            return ec2['instance'].state == 'running'
+
+        if waituntil(instance_running, timeout=120, interval=5):
+            info._ec2['connection'].deregister_image(info._ec2['image'])
+            info._ec2['snapshot'].delete()
+        else:
+            logger.error('Timeout while booting instance')