Implement roundrobin

Author: DifferentialOrange

tarantool/connection_pool.py

@@ -1,12 +1,15 @@
 # -*- coding: utf-8 -*-
 
 import json
+from enum import Enum
+from itertools import cycle
 
 from tarantool.connection import Connection, ConnectionInterface
 from tarantool.const import (CLUSTER_DISCOVERY_DELAY, CONNECTION_TIMEOUT,
                              RECONNECT_DELAY, RECONNECT_MAX_ATTEMPTS,
                              SOCKET_TIMEOUT)
-from tarantool.error import ClusterConnectWarning, ConfigurationError, warn
+from tarantool.error import (ClusterConnectWarning, ClusterTolopogyWarning,
+                             ConfigurationError, warn)
 from tarantool.utils import ENCODING_DEFAULT
 
 try:
@@ -69,6 +72,65 @@ def format_error(address, err):
     return format_error(address, 'port must be an int or a string')
 
 
+class RoundRobinStrategy(object):
+    """
+    Simple round-robin connection rotation
+    """
+    def __init__(self, pool):
+        self.RW_iter = None
+        self.RO_iter = None
+        self.default = None
+
+        if len(pool) == 0:
+            return
+        self.update(pool)
+
+    def update(self, new_pool):
+        self.default = next(iter(new_pool.values()))['conn']
+
+        RW_pool = []
+        RO_pool = []
+
+        for v in new_pool.values():
+            if v['ro'] is None:
+                continue
+
+            if v['ro']:
+                RO_pool.append(v['conn'])
+            else:
+                RW_pool.append(v['conn'])
+
+        if len(RW_pool) > 0:
+            self.RW_iter = cycle(RW_pool)
+
+        if len(RO_pool) > 0:
+            self.RO_iter = cycle(RO_pool)
+
+    def getnext(self, mode):
+        if mode == Mode.RW:
+            if self.RW_iter is not None:
+                return next(self.RW_iter)
+            else:
+                warn("can't find rw instance in pool, route request to any connection", ClusterTolopogyWarning)
+                return self.default
+        elif mode == Mode.PREFER_RO:
+            if self.RO_iter is not None:
+                return next(self.RO_iter)
+            elif self.RW_iter is not None:
+                return next(self.RW_iter)
+            else:
+                warn("can't find rw/ro instance in pool, route request to any connection", ClusterTolopogyWarning)
+                return self.default
+        elif mode == Mode.PREFER_RW:
+            if self.RW_iter is not None:
+                return next(self.RW_iter)
+            elif self.RO_iter is not None:
+                return next(self.RO_iter)
+            else:
+                warn("can't find rw/ro instance in pool, route request to any connection", ClusterTolopogyWarning)
+                return self.default
+
+
 class ConnectionPool(ConnectionInterface):
     def __init__(self,
                  addrs,
@@ -80,10 +142,10 @@ def __init__(self,
                  connect_now=True,
                  encoding=ENCODING_DEFAULT,
                  call_16=False,
-                 connection_timeout=CONNECTION_TIMEOUT,):
+                 connection_timeout=CONNECTION_TIMEOUT,
+                 strategy_class=RoundRobinStrategy,):
         if not isinstance(addrs, list) or len(addrs) == 0:
             raise ConfigurationError("addrs must be non-empty list")
-            addrs = []
 
         # Verify addresses.
         for addr in addrs:
@@ -94,21 +156,24 @@ def __init__(self,
 
         # Create connections
         self.pool = {}
+        self.strategy = strategy_class(self.pool)
 
         for addr in self.addrs:
             key = json.dumps(addr)
-            self.pool[key] = Connection(
-                host=addr['host'],
-                port=addr['port'],
-                user=user,
-                password=password,
-                socket_timeout=socket_timeout,
-                reconnect_max_attempts=reconnect_max_attempts,
-                reconnect_delay=reconnect_delay,
-                connect_now=False,
-                encoding=encoding,
-                call_16=call_16,
-                connection_timeout=connection_timeout)
+            self.pool[key] = {
+                'conn': Connection(
+                    host=addr['host'],
+                    port=addr['port'],
+                    user=user,
+                    password=password,
+                    socket_timeout=socket_timeout,
+                    reconnect_max_attempts=reconnect_max_attempts,
+                    reconnect_delay=reconnect_delay,
+                    connect_now=False, # Connect in ConnectionPool.connect()
+                    encoding=encoding,
+                    call_16=call_16,
+                    connection_timeout=connection_timeout)
+            }
 
         if connect_now:
             self.connect()
@@ -122,14 +187,27 @@ def is_closed(self):
     def connect(self):
         for addr in self.addrs:
             key = json.dumps(addr)
-            self.pool[key].connect()
+            self.pool[key]['conn'].connect()
 
-            if not self.pool[key].connected:
+            if not self.pool[key]['conn'].connected:
                 msg = 'Failed to connect to %s:%s' % str(addr['host'], addr['port'])
                 warn(msg, ClusterConnectWarning)
+                continue
 
-    def call(self, func_name, *args):
-        raise NotImplementedError
+            resp = self.pool[key]['conn'].call('box.info')
+            if not resp.data or type(resp.data[0]) != dict or resp.data[0]['ro'] is None:
+                warn("got incorrect response for rw/ro state, skipping instance", ClusterTolopogyWarning)
+                continue
+
+            self.pool[key]['ro'] = resp.data[0]['ro']
+
+        self.strategy.update(self.pool)
+
+    def call(self, func_name, *args, **kwargs):
+        mode = kwargs.get('mode', Mode.PREFER_RW)
+        conn = self.strategy.getnext(mode)
+        
+        return conn.call(func_name, args)
 
     def eval(self, expr, *args):
         raise NotImplementedError

tarantool/error.py

@@ -227,6 +227,11 @@ class ClusterConnectWarning(UserWarning):
     pass
 
 
+class ClusterTolopogyWarning(UserWarning):
+    '''Warning related to ro/rw cluster pool connection'''
+    pass
+
+
 # always print this warnings
 warnings.filterwarnings("always", category=NetworkWarning)