allow wait, wait_retries and wait_interval parameters for haproxy module. This allows the haproxy to wait for status "UP" when state=enabled and status "MAINT" when state=disabled

network/haproxy.py

@@ -68,6 +68,21 @@ options:
       - When disabling server, immediately terminate all the sessions attached to the specified server. This can be used to terminate long-running sessions after a server is put into maintenance mode, for instance.
     required: false
     default: false
+  wait:
+    description:
+      - Wait until the server reports a status of 'UP' when state=enabled, or status of 'MAINT' when state=disabled
+    required: false
+    default: false
+  wait_retries:
+    description:
+      - number of times to check for status after changing the state
+    required: false
+    default: 20
+  wait_interval:
+    description:
+      - number of seconds to wait between retries
+    required: false
+    default: 1
 '''
 
 EXAMPLES = '''
@@ -82,12 +97,21 @@ examples:
 # disable server, provide socket file
 - haproxy: state=disabled host={{ inventory_hostname }} socket=/var/run/haproxy.sock backend=www
 
+# disable server, provide socket file, wait until status reports in maintenance
+- haproxy: state=disabled host={{ inventory_hostname }} socket=/var/run/haproxy.sock backend=www wait=yes
+
 # disable backend server in 'www' backend pool and drop open sessions to it
 - haproxy: state=disabled host={{ inventory_hostname }} backend=www socket=/var/run/haproxy.sock shutdown_sessions=true
 
 # enable server in 'www' backend pool
 - haproxy: state=enabled host={{ inventory_hostname }} backend=www
 
+# enable server in 'www' backend pool wait until healthy
+- haproxy: state=enabled host={{ inventory_hostname }} backend=www wait=yes
+
+# enable server in 'www' backend pool wait until healthy. Retry 10 times with intervals of 5 seconds to retrieve the health
+- haproxy: state=enabled host={{ inventory_hostname }} backend=www wait=yes wait_retries=10 wait_interval=5
+
 # enable server in 'www' backend pool with change server(s) weight
 - haproxy: state=enabled host={{ inventory_hostname }} socket=/var/run/haproxy.sock weight=10 backend=www
 
@@ -95,11 +119,15 @@ author: "Ravi Bhure (@ravibhure)"
 '''
 
 import socket
+import csv
+import time
 
 
 DEFAULT_SOCKET_LOCATION="/var/run/haproxy.sock"
 RECV_SIZE = 1024
 ACTION_CHOICES = ['enabled', 'disabled']
+WAIT_RETRIES=20
+WAIT_INTERVAL=1
 
 ######################################################################
 class TimeoutException(Exception):
@@ -126,10 +154,12 @@ class HAProxy(object):
         self.weight = self.module.params['weight']
         self.socket = self.module.params['socket']
         self.shutdown_sessions = self.module.params['shutdown_sessions']
-
+        self.wait = self.module.params['wait']
+        self.wait_retries = self.module.params['wait_retries']
+        self.wait_interval = self.module.params['wait_interval']
         self.command_results = []
 
-    def execute(self, cmd, timeout=200):
+    def execute(self, cmd, timeout=200, capture_output=True):
         """
         Executes a HAProxy command by sending a message to a HAProxy's local
         UNIX socket and waiting up to 'timeout' milliseconds for the response.
@@ -144,10 +174,35 @@ class HAProxy(object):
         while buf:
             result += buf
             buf = self.client.recv(RECV_SIZE)
+        if capture_output:
             self.command_results = result.strip()
         self.client.close()
         return result
 
+    def wait_until_status(self, pxname, svname, status):
+        """
+        Wait for a server to become active (status == 'UP'). Try RETRIES times
+        with INTERVAL seconds of sleep in between. If the service has not reached
+        the expected status in that time, the module will fail. If the service was 
+        not found, the module will fail.
+        """
+        for i in range(1, self.wait_retries):
+            data = self.execute('show stat', 200, False).lstrip('# ')
+            r = csv.DictReader(data.splitlines())
+            found = False
+            for row in r:
+                if row['pxname'] == pxname and row['svname'] == svname:
+                    found = True
+                    if row['status'] == status:
+                        return True;
+                    else:
+                        time.sleep(self.wait_interval)
+
+            if not found:
+                self.module.fail_json(msg="unable to find server %s/%s" % (pxname, svname))
+
+        self.module.fail_json(msg="server %s/%s not status '%s' after %d retries. Aborting." % (pxname, svname, status, self.wait_retries))
+
     def enabled(self, host, backend, weight):
         """
         Enabled action, marks server to UP and checks are re-enabled,
@@ -170,6 +225,8 @@ class HAProxy(object):
                     if weight:
                         cmd += "; set weight %s/%s %s" % (pxname, svname, weight)
                     self.execute(cmd)
+                    if self.wait:
+                        self.wait_until_status(pxname, svname, 'UP')
 
         else:
             pxname = backend
@@ -177,6 +234,8 @@ class HAProxy(object):
             if weight:
                 cmd += "; set weight %s/%s %s" % (pxname, svname, weight)
             self.execute(cmd)
+            if self.wait:
+                self.wait_until_status(pxname, svname, 'UP')
 
     def disabled(self, host, backend, shutdown_sessions):
         """
@@ -200,6 +259,8 @@ class HAProxy(object):
                     if shutdown_sessions:
                         cmd += "; shutdown sessions server %s/%s" % (pxname, svname)
                     self.execute(cmd)
+                    if self.wait:
+                        self.wait_until_status(pxname, svname, 'MAINT')
 
         else:
             pxname = backend
@@ -207,6 +268,8 @@ class HAProxy(object):
             if shutdown_sessions:
                 cmd += "; shutdown sessions server %s/%s" % (pxname, svname)
             self.execute(cmd)
+            if self.wait:
+                self.wait_until_status(pxname, svname, 'MAINT')
 
     def act(self):
         """
@@ -236,6 +299,9 @@ def main():
             weight=dict(required=False, default=None),
             socket = dict(required=False, default=DEFAULT_SOCKET_LOCATION),
             shutdown_sessions=dict(required=False, default=False),
+            wait=dict(required=False, default=False),
+            wait_retries=dict(required=False, default=WAIT_RETRIES),
+            wait_interval=dict(required=False, default=WAIT_INTERVAL),
         ),
 
     )