From 294ef7d849b8227de52092ab2bcb28b3cdb24276 Mon Sep 17 00:00:00 2001
From: David Shrewsbury <shrewsbury.dave@gmail.com>
Date: Wed, 10 Oct 2012 16:50:41 -0400
Subject: [PATCH] Add set_algorithm() method to driver API.

HAProxy driver can now be given an algorithm to use for load
balancing.
---
 libra/worker/controller.py             | 21 +++++++++++++++++++++
 libra/worker/drivers/base.py           |  8 ++++++++
 libra/worker/drivers/haproxy/driver.py | 11 ++++++++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/libra/worker/controller.py b/libra/worker/controller.py
index d54cc6b0..ed7442c1 100644
--- a/libra/worker/controller.py
+++ b/libra/worker/controller.py
@@ -13,6 +13,7 @@
 # under the License.
 
 from libra.common.faults import BadRequest
+from libra.worker.drivers.base import LoadBalancerDriver
 
 
 class LBaaSController(object):
@@ -91,6 +92,26 @@ class LBaaSController(object):
                 self.msg[self.RESPONSE_FIELD] = self.RESPONSE_FAILURE
                 return self.msg
 
+        if 'algorithm' in self.msg:
+            algo = self.msg['algorithm'].upper()
+            if algo == 'ROUND_ROBIN':
+                algo = LoadBalancerDriver.ROUNDROBIN
+            elif algo == 'LEAST_CONNECTIONS':
+                algo = LoadBalancerDriver.LEASTCONN
+
+            try:
+                self.driver.set_algorithm(algo)
+            except NotImplementedError:
+                self.logger.error(
+                    "Selected driver does not support setting algorithm."
+                )
+                self.msg[self.RESPONSE_FIELD] = self.RESPONSE_FAILURE
+                return self.msg
+            except Exception as e:
+                self.logger.error("Selected driver failed setting algorithm.")
+                self.msg[self.RESPONSE_FIELD] = self.RESPONSE_FAILURE
+                return self.msg
+
         for lb_node in self.msg['nodes']:
             port, address = None, None
 
diff --git a/libra/worker/drivers/base.py b/libra/worker/drivers/base.py
index 13d4c0f8..7813cd23 100644
--- a/libra/worker/drivers/base.py
+++ b/libra/worker/drivers/base.py
@@ -30,6 +30,10 @@ class LoadBalancerDriver(object):
     via these API calls until the activate() method is called.
     """
 
+    # Load balancer algorithms
+    ROUNDROBIN = 1
+    LEASTCONN = 2
+
     def init(self):
         """ Allows the driver to do any initialization for a new config. """
         raise NotImplementedError()
@@ -42,6 +46,10 @@ class LoadBalancerDriver(object):
         """ Set the protocol of the instance. """
         raise NotImplementedError()
 
+    def set_algorithm(self, algo):
+        """ Set the algorithm used by the load balancer. """
+        raise NotImplementedError()
+
     def create(self):
         """ Create the load balancer. """
         raise NotImplementedError()
diff --git a/libra/worker/drivers/haproxy/driver.py b/libra/worker/drivers/haproxy/driver.py
index 8f72c663..2769444c 100644
--- a/libra/worker/drivers/haproxy/driver.py
+++ b/libra/worker/drivers/haproxy/driver.py
@@ -28,6 +28,7 @@ class HAProxyDriver(LoadBalancerDriver):
     def _init_config(self):
         self._config = dict()
         self.set_protocol('HTTP', 80)
+        self.set_algorithm(self.ROUNDROBIN)
 
     def _bind(self, address, port):
         self._config['bind_address'] = address
@@ -60,7 +61,7 @@ class HAProxyDriver(LoadBalancerDriver):
         output.append('    timeout connect 5000ms')
         output.append('    timeout client 50000ms')
         output.append('    timeout server 5000ms')
-        output.append('    balance roundrobin')
+        output.append('    balance %s' % self._config['algorithm'])
         output.append('    cookie SERVERID rewrite')
         output.append('frontend http-in')
         output.append('    bind %s:%s' % (self._config['bind_address'],
@@ -171,6 +172,14 @@ class HAProxyDriver(LoadBalancerDriver):
         else:
             self._bind('0.0.0.0', port)
 
+    def set_algorithm(self, algo):
+        if algo == self.ROUNDROBIN:
+            self._config['algorithm'] = 'roundrobin'
+        elif algo == self.LEASTCONN:
+            self._config['algorithm'] = 'leastconn'
+        else:
+            raise Exception('Invalid algorithm')
+
     def create(self):
         self._write_config()
         self._restart()