diff --git a/beanstalkc.py b/beanstalkc.py old mode 100644 new mode 100755 index d1527ef..8185a69 --- a/beanstalkc.py +++ b/beanstalkc.py @@ -17,11 +17,13 @@ limitations under the License. ''' -__version__ = '0.3.0' +__version__ = '0.4.0' import logging import socket - +import random +import time +import threading DEFAULT_HOST = 'localhost' DEFAULT_PORT = 11300 @@ -42,6 +44,187 @@ def wrap(wrapped_function, *args, **kwargs): except socket.error, err: raise SocketError(err) +class Pool(object): + """ Pool class enclosing multiple Connection-objects. The constructor + expects a list of tuples, (str(host),int(port)) as input.""" + def __init__(self,bstalks): + self.conLock = threading.RLock() + self.bstalks = bstalks + self.connections = [] + self._init_connections() + + logging.debug('DEBUG: ' + str(len(self.connections)) + ' out of ' + + str(len(bstalks)) + ' succeeded in connecting. (Yay?)') + + def _init_connections(self): + """Blocking until each connection has been tried. Any failures entails + a separate retry-thread.""" + for (host,port) in self.bstalks: + # Someone could have supplied the port as a string. Just convert it. + try: + port = int(port) + except ValueError: + logging.error('beanstalkc-pool failed converting %s to an int. Skipping connection.' % port) + continue + # Connect to the demon at host:port + try: + conn = Connection(host=host, port=port, parse_yaml=True, + connect_timeout=5) + with self.conLock: + self.connections.append( conn ) + except SocketError, e: + # Some network-error happened. Spawn a retry-thread: + logging.error('beanstalkc-pool failed connecting to %s %d. Retrying in a while.' % (host,port)) + t = threading.Thread(target=self._connect, args=(host,port)) + t.start() + + def _connect(self,host,port): + """Attempt connecting with beanstalkd at host:port. If successful add + to the connections-list and return. Else sleep, and try again.""" + sleep_time = 2 + while (True): + try: + conn = Connection(host=host, port=port, parse_yaml=True, + connect_timeout=5) + with self.conLock: + self.connections.append( conn ) + return + except SocketError, e: + logging.error('beanstalkc-pool failed connecting to %s %d. Retrying in a while.' % (host,port)) + time.sleep(sleep_time) + sleep_time += sleep_time % 64 + + def _call_wrap(self,conn,func,args): + """Handles the call of Connection-functions. Also responsible for + spawning threads to reconnect with lost Connections.""" + try: + if (args == None): + return func(conn) + else: + return func(conn, args ) + except SocketError as e: + with self.conLock: + self.connections.remove(conn) + t = threading.Thread(target=self._connect, args=(conn.host,conn.port)) + t.start() + logging.error('beanstalkc-pool socket-error to %s %d. Started reconnect-thread.' % (conn.host,conn.port)) + raise e + + def _send_to_rand_conn(self,func,args=None): + """Send to a random connection in the pool. Returns a tuple of the + connection the command was sent to, and its response.""" + retries = 5 + while (True): + try: + with self.conLock: + conn = random.choice(self.connections) + break + except IndexError as e: + # self.connections is empty.. Sleep for a little while and try + # again. + if (not retries): + raise SocketError('No connections available.') + time.sleep(2) + retries -= 1 + + resp = self._call_wrap(conn,func,args) + return (conn,resp) + + def _send_to_all(self,func,args=None): + """Send to all connections in the pool. Returns a list of tuples + (conn,response)""" + results = [] + with self.conLock: + # Loop over a copy of self.connections (slice notation) to + # encompass changes to itself, during the loop. + if len(self.connections) == 0: + raise BeanstalkcException('Pool is empty. Nothing sent.') + for conn in self.connections[:]: + try: + results.append( (conn,self._call_wrap(conn,func,args)) ) + except SocketError as e: + logging.error('beanstalkc-pool connection error in _send_to_all(). Skipping connection.') + return results + + def reserve(self,pool_timeout=None,timeout=10): + """Reserves a job from the pool randomly. + pool_timeout: Overall time in seconds before this call returns. + timeout: Individual reserve-calls timeouts. + + Returns: (connection,result) or (None,None)""" + timer = None + if (pool_timeout != None): + timer = time.time() + while (True): + try: + with self.conLock: + conn = random.choice(self.connections) + result = self._call_wrap(conn, Connection.reserve, timeout) + # If result == None, the reserve timed-out and we try again. + if (result != None): + return (conn,result) + except DeadlineSoon as e: + # Deadlinesoon. We give up and try reserving another job. + pass + except SocketError as e: + # The connection we chose failed. Trying another.. (it has been + # removed from the connections list..) + pass + except IndexError as e: + # self.connections is empty.. Sleep for a little while and try again. + time.sleep(2) + if (timer != None and time.time() > timer + pool_timeout): + # If we get here, the pool-timeout has been reached. Returning + # None just as Connection.reserve when the timeout is reached. + return (None,None) + + def close(self): + """Close all connections in the pool.""" + self._send_to_all( Connection.close) + + def put(self,arg): + """Put a job in the pool.""" + return self._send_to_rand_conn( Connection.put, arg) + + def tubes(self): + """Return a list of lists of all existing tubes.""" + return self._send_to_all( Connection.tubes) + + def using(self): + """Return a list of (connection,tube) indicating the tube currently in + use for every connection in the pool.""" + return self._send_to_all( Connection.use) + + def use(self,name): + """Use 'name' tube on every connection in the pool.""" + self._send_to_all( Connection.use, name) + + def watching(self): + """Return a list of (connection,tubes) where tubes is all the tubes + being watched in that conenction.""" + return self._send_to_all( Connection.watching) + + def watch(self, name): + """Watch a given tube in all connections.""" + return self._send_to_all( Connection.watch, name) + + def ignore(self,name): + """Ignore the given tube in all connections in the pool.""" + self._send_to_all( Connection.ignore, name) + + def stats(self): + """Return a list of (connection,dicts) with beanstalkd statistics.""" + return self._send_to_all( Connection.stats) + + def stats_tube(self, name): + """Return a list of dicts of stats about a given tube.""" + return self._send_to_all( Connection.stats_tube, name) + + def pause_tube(self, name, delay): + """Pause a tube on all connections for a given delay time, in + seconds.""" + self._send_to_all( Connection.pause_tube,[name, delay]) + class Connection(object): def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, parse_yaml=True, @@ -92,7 +275,7 @@ def _interact(self, command, expected_ok, expected_err=[]): def _read_response(self): line = SocketError.wrap(self._socket_file.readline) if not line: - raise SocketError() + raise SocketError('SocketError') response = line.split() return response[0], response[1:]