Source code for zktools.connection

# ***** BEGIN LICENSE BLOCK *****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
#
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
#
# The Original Code is Mozilla zktools.
#
# The Initial Developer of the Original Code is Mozilla Foundation.
# Portions created by the Initial Developer are Copyright (C) 2011
# the Initial Developer. All Rights Reserved.
#
# Contributor(s):
# Ben Bangert (bbangert@mozilla.com)
#
# Alternatively, the contents of this file may be used under the terms of
# either the GNU General Public License Version 2 or later (the "GPL"), or
# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
#
# ***** END LICENSE BLOCK *****
"""Zookeeper Connection Classes"""
import time
import threading

import zookeeper


[docs]class ZkConnection(object): """Zookeeper Connection object"""
[docs] def __init__(self, host="localhost:2181", connect_timeout=10, session_timeout=10 * 1000, reconnect_timeout=10 * 1000): """Create a connection object, capable of automatically reconnecting when the connection is lost. In the event the connection is lost, the zookeeper command will block until the connection is available again, or the ``reconnect_timeout`` is reached. The latter will raise an Exception about the reconnect period expiring. Example:: conn = ZkConnection() conn.connect() node = conn.create( "/my/node", "a value", [{"perms": 0x1f, "scheme": "world", "id": "anyone"}], 0) .. note:: ZkConnection ensures that all zookeeper functions use the same zookeeper handle that ZkConnection has, so it does not need to be passed in as the first argument. :param host: A valid zookeeper host string :param connect_timeout: Timeout for connecting to zookeeper :type connect_timeout: int :param session_timeout: Timeout for the zookeeper session :type session_timeout: int :param reconnect_timeout: How many seconds to attempt to reconnect before giving up. Set to 0 to avoid reconnect attempts. :type reconnect_timeout: int """ self.connected = False self._host = host self._connect_timeout = connect_timeout self._session_timeout = session_timeout self._reconnect_timeout = reconnect_timeout self._cv = threading.Condition() self._handle = None
def _handle_connection(self, handle, typ, state, path): # The Zookeeper API runs this in a separate event thread with self._cv: if typ == zookeeper.SESSION_EVENT: if state == zookeeper.CONNECTED_STATE: self.connected = True elif state == zookeeper.CONNECTING_STATE: self.connected = False elif state in (zookeeper.EXPIRED_SESSION_STATE, zookeeper.AUTH_FAILED_STATE): # Last event for this connection, session is dead, # clean up self._handle = None self.connected = False self._cv.notify_all()
[docs] def connect(self): """Connect to zookeeper""" # Run this with our threading Condition with self._cv: # First, check that we didn't just connect if self.connected: return # See if the client is already attempting a connection if self._handle is not None: start_time = time.time() time_taken = 0 while time_taken <= self._reconnect_timeout: # Release and wait until we hit our timeout # NOTE: We loop here, because the connection handler # releases for *every state change*, and we only care # about getting connected self._cv.wait(self._reconnect_timeout - time_taken) if self.connected: return # Triggered, if the session is expired, we break out # to let logic continue below if self._handle is None: break time_taken = time.time() - start_time if self._handle is not None and not self.connected: raise Exception("Timed out waiting for reconnect.") # Either first run, or a prior session was ditched entirely self._handle = zookeeper.init(self._host, self._handle_connection, self._session_timeout) self._cv.wait(self._connect_timeout) if not self.connected: raise Exception("Unable to connect to Zookeeper")
[docs] def __getattr__(self, name): """Returns a reconnecting version that also uses the current handle""" zoo_func = getattr(zookeeper, name) # Check that we're still connected if not self.connected: self.connect() def call_func(*args, **kwargs): # We wait/try this until we're connected, unless it took # too long start_time = time.time() time_taken = 0 while time_taken <= self._reconnect_timeout: try: return zoo_func(self._handle, *args, **kwargs) except (zookeeper.ConnectionLossException, zookeeper.SessionExpiredException, zookeeper.SessionMovedException): self.connect() except zookeeper.ZooKeeperException, msg: if 'zhandle already freed' in msg: self._handle = None self.connected = False self.connect() else: raise time_taken = time.time() - start_time raise Exception("Unable to reconnect to execute command.") call_func.__doc__ = zoo_func.__doc__ # Set this function on ourself so that further calls bypass # getattr setattr(self, name, call_func) return call_func
def __dir__(self): return self.__dict__.keys() + zookeeper.__dict__.keys()