From 899811f339ee585997ab734793651264434f0b74 Mon Sep 17 00:00:00 2001
From: Bruce Pennypacker <bpennypacker@care.com>
Date: Fri, 14 Feb 2014 20:16:58 +0000
Subject: [PATCH] Added wait_timeout parameter

---
 library/cloud/ec2_elb | 29 ++++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/library/cloud/ec2_elb b/library/cloud/ec2_elb
index c6f4a72b0e..fe5260bb16 100644
--- a/library/cloud/ec2_elb
+++ b/library/cloud/ec2_elb
@@ -82,6 +82,12 @@ options:
     choices: ["yes", "no"]
     aliases: []
     version_added: "1.5"
+  wait_timeout:
+    description:
+      - Number of seconds to wait for an instance to change state. If 0 then this module may return an error if a transient error occurs. If non-zero then any transient errors are ignored until the timeout is reached. Ignored when wait=no.
+    required: false
+    default: 0
+    version_added: "1.5"
 
 """
 
@@ -133,7 +139,7 @@ class ElbManager:
         self.lbs = self._get_instance_lbs(ec2_elbs)
         self.changed = False
 
-    def deregister(self, wait):
+    def deregister(self, wait, timeout):
         """De-register the instance from all ELBs and wait for the ELB
         to report it out-of-service"""
 
@@ -146,13 +152,13 @@ class ElbManager:
                 return
 
             if wait:
-                self._await_elb_instance_state(lb, 'OutOfService', initial_state)
+                self._await_elb_instance_state(lb, 'OutOfService', initial_state, timeout)
             else:
                 # We cannot assume no change was made if we don't wait
                 # to find out
                 self.changed = True
 
-    def register(self, wait, enable_availability_zone):
+    def register(self, wait, enable_availability_zone, timeout):
         """Register the instance for all ELBs and wait for the ELB
         to report the instance in-service"""
         for lb in self.lbs:
@@ -165,7 +171,7 @@ class ElbManager:
             lb.register_instances([self.instance_id])
 
             if wait:
-                self._await_elb_instance_state(lb, 'InService', initial_state)
+                self._await_elb_instance_state(lb, 'InService', initial_state, timeout)
             else:
                 # We cannot assume no change was made if we don't wait
                 # to find out
@@ -195,10 +201,12 @@ class ElbManager:
         # lb.availability_zones
         return instance.placement in lb.availability_zones
 
-    def _await_elb_instance_state(self, lb, awaited_state, initial_state):
+    def _await_elb_instance_state(self, lb, awaited_state, initial_state, timeout):
         """Wait for an ELB to change state
         lb: load balancer
         awaited_state : state to poll for (string)"""
+
+        wait_timeout = time.time() + timeout
         while True:
             instance_state = self._get_instance_health(lb)
 
@@ -217,7 +225,8 @@ class ElbManager:
                 # If it's pending, we'll skip further checks andd continue waiting
                 pass
             elif (awaited_state == 'InService'
-                  and instance_state.reason_code == "Instance"):
+                  and instance_state.reason_code == "Instance"
+                  and time.time() >= wait_timeout):
                 # If the reason_code for the instance being out of service is
                 # "Instance" this indicates a failure state, e.g. the instance
                 # has failed a health check or the ELB does not have the
@@ -303,7 +312,8 @@ def main():
             ec2_access_key={'default': None, 'aliases': ['aws_access_key', 'access_key']},
             region={'default': None, 'required': False, 'aliases':['aws_region', 'ec2_region'], 'choices':AWS_REGIONS},
             enable_availability_zone={'default': True, 'required': False, 'choices': BOOLEANS, 'type': 'bool'},
-            wait={'required': False, 'choices': BOOLEANS, 'default': True, 'type': 'bool'}
+            wait={'required': False, 'choices': BOOLEANS, 'default': True, 'type': 'bool'},
+            wait_timeout={'requred': False, 'default': 0, 'type': 'int'}
         )
     )
 
@@ -315,6 +325,7 @@ def main():
     region = module.params['region']
     wait = module.params['wait']
     enable_availability_zone = module.params['enable_availability_zone']
+    timeout = module.params['wait_timeout']
 
     if module.params['state'] == 'present' and 'ec2_elbs' not in module.params:
         module.fail_json(msg="ELBs are required for registration")
@@ -330,9 +341,9 @@ def main():
                 module.fail_json(msg=msg)
 
     if module.params['state'] == 'present':
-        elb_man.register(wait, enable_availability_zone)
+        elb_man.register(wait, enable_availability_zone, timeout)
     elif module.params['state'] == 'absent':
-        elb_man.deregister(wait)
+        elb_man.deregister(wait, timeout)
 
     ansible_facts = {'ec2_elbs': [lb.name for lb in elb_man.lbs]}
     ec2_facts_result = dict(changed=elb_man.changed, ansible_facts=ansible_facts)