def test_canceled_join_long_time(self): zk = ZooKeeper(self._server.ensemble) session_id = zk.session_id() ss = ServerSet(zk, self.SERVICE_PATH) join_signal = threading.Event() memberships = [] def on_expire(): pass def do_join(): memberships.append(ss.join(self.INSTANCE1, expire_callback=on_expire)) class JoinThread(threading.Thread): def run(_): while True: join_signal.wait() join_signal.clear() do_join() joiner = JoinThread() joiner.daemon = True joiner.start() do_join() assert len(memberships) == 1 and memberships[0] is not Membership.error() self._server.expire(session_id) self._server.shutdown() join_signal.set() self._server.start() while len(memberships) == 1: time.sleep(0.1) assert len(memberships) == 2 and memberships[1] is not Membership.error()
def test_session_event(): with ZookeeperServer() as server: disconnected = threading.Event() def on_event(zk, event, state, _): if zk._live.is_set() and state != zookeeper.CONNECTED_STATE: disconnected.set() zk = ZooKeeper(server.ensemble, watch=on_event) session_id = zk.session_id() children = [] completion_event = threading.Event() def children_completion(_, rc, results): children.extend(results) completion_event.set() server.shutdown() disconnected.wait(timeout=MAX_EVENT_WAIT_SECS) assert disconnected.is_set() zk.aget_children('/', None, children_completion) # expire session server.expire(session_id) server.start() completion_event.wait(timeout=MAX_EVENT_WAIT_SECS) assert completion_event.is_set() assert children == ['zookeeper']
def test_shard_id_registers(self): ss1 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) ss2 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) ss1.join(self.INSTANCE1, shard=0) ss2.join(self.INSTANCE2, shard=1) assert list(ss1) == [ ServiceInstance(self.INSTANCE1, shard=0), ServiceInstance(self.INSTANCE2, shard=1) ] assert list(ss2) == [ ServiceInstance(self.INSTANCE1, shard=0), ServiceInstance(self.INSTANCE2, shard=1) ]
def test_expand_ensemble(): m = mox.Mox() m.StubOutWithMock(socket, 'gethostbyname_ex') socket.gethostbyname_ex('localhost').AndReturn(('localhost', [], ['foo'])) socket.gethostbyname_ex('localhost').AndReturn(('localhost', [], ['bar'])) socket.gethostbyname_ex('localhost').AndReturn(('localhost', [], ['baz', 'bak'])) m.ReplayAll() assert ZooKeeper.expand_ensemble('localhost:1234') == 'foo:1234' assert ZooKeeper.expand_ensemble('localhost:1234,localhost') == 'bar:1234,baz:2181,bak:2181' m.UnsetStubs() m.VerifyAll()
def test_client_reconnect(): cluster = ZookeeperClusterBootstrapper() port = cluster.start(1) zk = ZooKeeper('localhost:%d' % port) zk.get_children('/') cluster.stop(1) with pytest.raises(zookeeper.ConnectionLossException): zk.get_children('/') cluster.start(1) with pytest.raises(zookeeper.ConnectionLossException): zk.get_children('/') zk.reconnect() assert zk.get_children('/') == ['zookeeper']
def test_async_client_iteration(self): ss1 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) ss2 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) ss1.join(self.INSTANCE1) ss2.join(self.INSTANCE2) assert list(ss1) == [ ServiceInstance(self.INSTANCE1), ServiceInstance(self.INSTANCE2) ] assert list(ss2) == [ ServiceInstance(self.INSTANCE1), ServiceInstance(self.INSTANCE2) ]
def test_expand_ensemble(): m = mox.Mox() m.StubOutWithMock(socket, 'gethostbyname_ex') socket.gethostbyname_ex('localhost').AndReturn(('localhost', [], ['foo'])) socket.gethostbyname_ex('localhost').AndReturn(('localhost', [], ['bar'])) socket.gethostbyname_ex('localhost').AndReturn( ('localhost', [], ['baz', 'bak'])) m.ReplayAll() assert ZooKeeper.expand_ensemble('localhost:1234') == 'foo:1234' assert ZooKeeper.expand_ensemble( 'localhost:1234,localhost') == 'bar:1234,baz:2181,bak:2181' m.UnsetStubs() m.VerifyAll()
def test_metrics(): with ZookeeperServer() as server: event = threading.Event() def watch_set(*args): event.set() zk = ZooKeeper(server.ensemble, watch=watch_set) zk._live.wait(timeout=MAX_EVENT_WAIT_SECS) sample = zk.metrics.sample() assert sample['live'] == 1 assert sample['session_id'] == zk.session_id assert sample['session_expirations'] == 0 assert sample['connection_losses'] == 0 old_session_id = zk.session_id event.clear() server.expire(zk.session_id) event.wait(timeout=MAX_EXPIRE_WAIT_SECS) zk._live.wait(timeout=MAX_EVENT_WAIT_SECS) sample = zk.metrics.sample() assert sample['live'] == 1 assert sample['session_id'] == zk.session_id assert old_session_id != zk.session_id assert sample['session_expirations'] == 1
def test_client_connect_times_out(): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('localhost', 0)) _, port = sock.getsockname() sock.close() with pytest.raises(ZooKeeper.ConnectionTimeout): ZooKeeper('localhost:%d' % port, timeout_secs=1.0)
def test_safe_operations(): with ZookeeperServer() as server: zk = ZooKeeper(server.ensemble) assert zk.safe_create('/a/b/c/d') == '/a/b/c/d' session_id = zk.session_id() finish_event = threading.Event() class CreateThread(threading.Thread): def run(self): zk.safe_create('/foo/bar/baz/bak') finish_event.set() server.shutdown() server.expire(session_id) ct = CreateThread() ct.start() server.start() finish_event.wait(timeout=MAX_EXPIRE_WAIT_SECS) assert finish_event.is_set() assert zk.exists('/a/b/c/d') assert zk.exists('/foo/bar/baz/bak') session_id = zk.session_id() assert zk.safe_delete('/a') delete_event = threading.Event() class DeleteThread(threading.Thread): def run(self): zk.safe_delete('/foo') delete_event.set() server.shutdown() server.expire(session_id) dt = DeleteThread() dt.start() server.start() delete_event.wait(timeout=MAX_EXPIRE_WAIT_SECS) assert delete_event.is_set() assert not zk.exists('/a') assert not zk.exists('/foo')
def test_stopped(): with ZookeeperServer() as server: zk = ZooKeeper('localhost:%d' % server.zookeeper_port) assert zk.get_children('/') == ['zookeeper'] zk.stop() with pytest.raises(ZooKeeper.Stopped): zk.get_children('/')
def test_client_iteration(self): ss = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) assert list(ss) == [] ss.join(self.INSTANCE1) assert list(ss) == [ServiceInstance(self.INSTANCE1)] ss.join(self.INSTANCE2) assert list(ss) == [ ServiceInstance(self.INSTANCE1), ServiceInstance(self.INSTANCE2) ]
def _construct_serverset(self, options): import socket import threading import zookeeper from twitter.common.zookeeper.client import ZooKeeper from twitter.common.zookeeper.serverset import Endpoint, ServerSet log.debug('ServerSet module constructing serverset.') hostname = socket.gethostname() primary_port = int(options.serverset_module_primary_port) primary = Endpoint(hostname, primary_port) additional = dict((port_name, Endpoint(hostname, port_number)) for port_name, port_number in options.serverset_module_extra.items()) # TODO(wickman) Add timeout parameterization here. self._zookeeper = ZooKeeper(options.serverset_module_ensemble) self._serverset = ServerSet(self._zookeeper, options.serverset_module_path) self._join_args = (primary, additional)
def test_client_watcher(self): canceled_endpoints = [] canceled = threading.Event() joined_endpoints = [] joined = threading.Event() def on_join(endpoint): joined_endpoints[:] = [endpoint] joined.set() def on_leave(endpoint): canceled_endpoints[:] = [endpoint] canceled.set() service1 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH, on_join=on_join, on_leave=on_leave) service2 = ServerSet(ZooKeeper(self._server.ensemble), self.SERVICE_PATH) member1 = service2.join(self.INSTANCE1) joined.wait(2.0) assert joined.is_set() assert not canceled.is_set() assert joined_endpoints == [ServiceInstance(self.INSTANCE1)] joined.clear() service2.join(self.INSTANCE2) joined.wait(2.0) assert joined.is_set() assert not canceled.is_set() assert joined_endpoints == [ServiceInstance(self.INSTANCE2)] joined.clear() service2.cancel(member1) canceled.wait(2.0) assert canceled.is_set() assert not joined.is_set() assert canceled_endpoints == [ServiceInstance(self.INSTANCE1)] canceled.clear()
def test_client_stops_propagate_through_completions(): with ZookeeperServer() as server: zk = ZooKeeper('localhost:%d' % server.zookeeper_port) server.shutdown() # zk.get_children should block until reconnected stopped_event = threading.Event() class GetThread(threading.Thread): def run(self): try: zk.get_children('/') except ZooKeeper.Stopped: stopped_event.set() gt = GetThread() gt.start() time.sleep(0.1) # guarantee an interpreter thread yield zk.stop() stopped_event.wait(timeout=MAX_EVENT_WAIT_SECS) assert stopped_event.is_set()
def main(args): if len(args) != 1: app.error('Must supply a serverset path to monitor.') def on_join(endpoint): print('@ %s += %s' % (datetime.now(), endpoint)) def on_leave(endpoint): print('@ %s -= %s' % (datetime.now(), endpoint)) ss = ServerSet(ZooKeeper(), args[0], on_join=on_join, on_leave=on_leave) while True: time.sleep(100)
def test_canceled_join_long_time(self): zk = ZooKeeper(self._server.ensemble) session_id = zk.session_id() ss = ServerSet(zk, self.SERVICE_PATH) join_signal = threading.Event() memberships = [] def on_expire(): pass def do_join(): memberships.append( ss.join(self.INSTANCE1, expire_callback=on_expire)) class JoinThread(threading.Thread): def run(_): while True: join_signal.wait() join_signal.clear() do_join() joiner = JoinThread() joiner.daemon = True joiner.start() do_join() assert len( memberships) == 1 and memberships[0] is not Membership.error() self._server.expire(session_id) self._server.shutdown() join_signal.set() self._server.start() while len(memberships) == 1: time.sleep(0.1) assert len( memberships) == 2 and memberships[1] is not Membership.error()
def _construct_serverset(self, options): import socket import threading import zookeeper from twitter.common.zookeeper.client import ZooKeeper from twitter.common.zookeeper.serverset import Endpoint, ServerSet log.debug("ServerSet module constructing serverset.") hostname = socket.gethostname() primary_port = int(options.serverset_module_primary_port) primary = Endpoint(hostname, primary_port) additional = dict( (port_name, Endpoint(hostname, port_number)) for port_name, port_number in options.serverset_module_extra.items() ) # TODO(wickman) Add timeout parameterization here. self._zookeeper = ZooKeeper(options.serverset_module_ensemble) self._serverset = ServerSet(self._zookeeper, options.serverset_module_path) self._join_args = (primary, additional) self._join_kwargs = {"shard": options.serverset_module_shard_id} if options.serverset_module_shard_id else {}
def test_bad_ensemble(): with pytest.raises(ZooKeeper.InvalidEnsemble): ZooKeeper.expand_ensemble('localhost:') with pytest.raises(ZooKeeper.InvalidEnsemble): ZooKeeper.expand_ensemble('localhost:sheeps') m = mox.Mox() m.StubOutWithMock(socket, 'gethostbyname_ex') socket.gethostbyname_ex('zookeeper.twitter.com').AndRaise( socket.gaierror(8, 'nodename nor servname provided, or not known')) m.ReplayAll() with pytest.raises(ZooKeeper.InvalidEnsemble): ZooKeeper.expand_ensemble('zookeeper.twitter.com:2181') m.UnsetStubs() m.VerifyAll()
class ServerSetModule(app.Module): """ Binds this application to a Zookeeper ServerSet. """ OPTIONS = { 'serverset-enable': options.Option( '--serverset-enable', default=False, action='store_true', dest='serverset_module_enable', help= 'Enable the ServerSet module. Requires --serverset-path and --serverset-primary.' ), 'serverset-ensemble': options.Option( '--serverset-ensemble', default='zookeeper.local.twitter.com:2181', dest='serverset_module_ensemble', metavar='HOST[:PORT]', help= 'The serverset ensemble to talk to. HOST or HOST:PORT pair. If the HOST is a RR DNS ' 'record, we fan out to the entire ensemble. If no port is specified, 2181 assumed.' ), 'serverset-path': options.Option( '--serverset-path', default=None, dest='serverset_module_path', metavar='PATH', type='str', help= 'The serverset path to join, preferably /twitter/service/(role)/(service)/(env) ' 'where env is prod, staging, devel.'), 'serverset-primary': options.Option('--serverset-primary', type='int', metavar='PORT', dest='serverset_module_primary_port', default=None, help='Port on which to bind the primary endpoint.'), 'serverset-shard-id': options.Option('--serverset-shard-id', type='int', metavar='INT', dest='serverset_module_shard_id', default=None, help='Shard id to assign this serverset entry.'), 'serverset-extra': options.Option( '--serverset-extra', default={}, type='string', nargs=1, action='callback', metavar='NAME:PORT', callback=add_port_to('serverset_module_extra'), dest='serverset_module_extra', help= 'Additional endpoints to bind. Format NAME:PORT. May be specified multiple times.' ), 'serverset-persistence': options.Option( '--serverset-persistence', '--no-serverset-persistence', action='callback', callback=set_bool, dest='serverset_module_persistence', default=True, help= 'If serverset persistence is enabled, if the serverset connection is dropped for any ' 'reason, we will retry to connect forever. If serverset persistence is turned off, ' 'the application will commit seppuku -- sys.exit(1) -- upon session disconnection.' ), } def __init__(self): app.Module.__init__(self, __name__, description="ServerSet module") self._zookeeper = None self._serverset = None self._membership = None self._join_args = None self._torndown = False self._rejoin_event = threading.Event() self._joiner = None @property def serverset(self): return self._serverset @property def zh(self): if self._zookeeper: return self._zookeeper._zh def _assert_valid_inputs(self, options): if not options.serverset_module_enable: return assert options.serverset_module_path is not None, ( 'If serverset module enabled, serverset path must be specified.') assert options.serverset_module_primary_port is not None, ( 'If serverset module enabled, serverset primary port must be specified.' ) assert isinstance( options.serverset_module_extra, dict), ('Serverset additional endpoints must be a dictionary!') for name, value in options.serverset_module_extra.items(): assert isinstance( name, str), 'Additional endpoints must be named by strings!' assert isinstance( value, int), 'Additional endpoint ports must be integers!' try: primary_port = int(options.serverset_module_primary_port) except ValueError as e: raise ValueError('Could not parse serverset primary port: %s' % e) def _construct_serverset(self, options): import socket import threading import zookeeper from twitter.common.zookeeper.client import ZooKeeper from twitter.common.zookeeper.serverset import Endpoint, ServerSet log.debug('ServerSet module constructing serverset.') hostname = socket.gethostname() primary_port = int(options.serverset_module_primary_port) primary = Endpoint(hostname, primary_port) additional = dict((port_name, Endpoint(hostname, port_number)) for port_name, port_number in options.serverset_module_extra.items()) # TODO(wickman) Add timeout parameterization here. self._zookeeper = ZooKeeper(options.serverset_module_ensemble) self._serverset = ServerSet(self._zookeeper, options.serverset_module_path) self._join_args = (primary, additional) self._join_kwargs = ({ 'shard': options.serverset_module_shard_id } if options.serverset_module_shard_id else {}) def _join(self): log.debug('ServerSet module joining serverset.') primary, additional = self._join_args self._membership = self._serverset.join( primary, additional, expire_callback=self.on_expiration, **self._join_kwargs) def on_expiration(self): if self._torndown: return log.debug('Serverset session expired.') if not app.get_options().serverset_module_persistence: log.debug('Committing seppuku...') sys.exit(1) else: log.debug('Rejoining...') self._rejoin_event.set() def setup_function(self): options = app.get_options() if options.serverset_module_enable: self._assert_valid_inputs(options) self._construct_serverset(options) self._thread = ServerSetJoinThread(self._rejoin_event, self._join) self._thread.start() self._rejoin_event.set() def teardown_function(self): self._torndown = True if self._membership: self._serverset.cancel(self._membership) self._zookeeper.stop()
def test_safe_create(): with ZookeeperServer() as server: zk_auth = ZooKeeper(server.ensemble, authentication=('digest', 'jack:jill')) zk_auth.safe_create('/a', acl=ZooDefs.Acls.EVERYONE_READ_CREATOR_ALL) assert zk_auth.exists('/a') zk_noauth = ZooKeeper(server.ensemble) with pytest.raises(zookeeper.NoAuthException): zk_noauth.safe_create('/a/b') assert not zk_auth.exists('/a/b') zk_auth.safe_create('/a/b', acl=ZooDefs.Acls.OPEN_ACL_UNSAFE) assert zk_noauth.exists('/a/b') zk_noauth.safe_create('/a/b/c') assert zk_noauth.exists('/a/b/c')
def test_expand_ensemble(): assert ZooKeeper.expand_ensemble('localhost:1234') == '127.0.0.1:1234'
class ServerSetModule(app.Module): """ Binds this application to a Zookeeper ServerSet. """ OPTIONS = { "serverset-enable": options.Option( "--serverset-enable", default=False, action="store_true", dest="serverset_module_enable", help="Enable the ServerSet module. Requires --serverset-path and --serverset-primary.", ), "serverset-ensemble": options.Option( "--serverset-ensemble", default="zookeeper.local.twitter.com:2181", dest="serverset_module_ensemble", metavar="HOST[:PORT]", help="The serverset ensemble to talk to. HOST or HOST:PORT pair. If the HOST is a RR DNS " "record, we fan out to the entire ensemble. If no port is specified, 2181 assumed.", ), "serverset-path": options.Option( "--serverset-path", default=None, dest="serverset_module_path", metavar="PATH", type="str", help="The serverset path to join, preferably /twitter/service/(role)/(service)/(env) " "where env is prod, staging, devel.", ), "serverset-primary": options.Option( "--serverset-primary", type="int", metavar="PORT", dest="serverset_module_primary_port", default=None, help="Port on which to bind the primary endpoint.", ), "serverset-shard-id": options.Option( "--serverset-shard-id", type="int", metavar="INT", dest="serverset_module_shard_id", default=None, help="Shard id to assign this serverset entry.", ), "serverset-extra": options.Option( "--serverset-extra", default={}, type="string", nargs=1, action="callback", metavar="NAME:PORT", callback=add_port_to("serverset_module_extra"), dest="serverset_module_extra", help="Additional endpoints to bind. Format NAME:PORT. May be specified multiple times.", ), "serverset-persistence": options.Option( "--serverset-persistence", "--no-serverset-persistence", action="callback", callback=set_bool, dest="serverset_module_persistence", default=True, help="If serverset persistence is enabled, if the serverset connection is dropped for any " "reason, we will retry to connect forever. If serverset persistence is turned off, " "the application will commit seppuku -- sys.exit(1) -- upon session disconnection.", ), } def __init__(self): app.Module.__init__(self, __name__, description="ServerSet module") self._zookeeper = None self._serverset = None self._membership = None self._join_args = None self._torndown = False self._rejoin_event = threading.Event() self._joiner = None @property def serverset(self): return self._serverset @property def zh(self): if self._zookeeper: return self._zookeeper._zh def _assert_valid_inputs(self, options): if not options.serverset_module_enable: return assert ( options.serverset_module_path is not None ), "If serverset module enabled, serverset path must be specified." assert ( options.serverset_module_primary_port is not None ), "If serverset module enabled, serverset primary port must be specified." assert isinstance(options.serverset_module_extra, dict), "Serverset additional endpoints must be a dictionary!" for name, value in options.serverset_module_extra.items(): assert isinstance(name, str), "Additional endpoints must be named by strings!" assert isinstance(value, int), "Additional endpoint ports must be integers!" try: primary_port = int(options.serverset_module_primary_port) except ValueError as e: raise ValueError("Could not parse serverset primary port: %s" % e) def _construct_serverset(self, options): import socket import threading import zookeeper from twitter.common.zookeeper.client import ZooKeeper from twitter.common.zookeeper.serverset import Endpoint, ServerSet log.debug("ServerSet module constructing serverset.") hostname = socket.gethostname() primary_port = int(options.serverset_module_primary_port) primary = Endpoint(hostname, primary_port) additional = dict( (port_name, Endpoint(hostname, port_number)) for port_name, port_number in options.serverset_module_extra.items() ) # TODO(wickman) Add timeout parameterization here. self._zookeeper = ZooKeeper(options.serverset_module_ensemble) self._serverset = ServerSet(self._zookeeper, options.serverset_module_path) self._join_args = (primary, additional) self._join_kwargs = {"shard": options.serverset_module_shard_id} if options.serverset_module_shard_id else {} def _join(self): log.debug("ServerSet module joining serverset.") primary, additional = self._join_args self._membership = self._serverset.join( primary, additional, expire_callback=self.on_expiration, **self._join_kwargs ) def on_expiration(self): if self._torndown: return log.debug("Serverset session expired.") if not app.get_options().serverset_module_persistence: log.debug("Committing seppuku...") sys.exit(1) else: log.debug("Rejoining...") self._rejoin_event.set() def setup_function(self): options = app.get_options() if options.serverset_module_enable: self._assert_valid_inputs(options) self._construct_serverset(options) self._thread = ServerSetJoinThread(self._rejoin_event, self._join) self._thread.start() self._rejoin_event.set() def teardown_function(self): self._torndown = True if self._membership: self._serverset.cancel(self._membership) self._zookeeper.stop()
def make_zk(server, **kw): return ZooKeeper('localhost:%d' % server.zookeeper_port, timeout_secs=CONNECT_TIMEOUT_SECS, max_reconnects=CONNECT_RETRIES, **kw)
def make_zk(cls, ensemble): return ZooKeeper(ensemble)
def make_zk(cls, ensemble, **kw): return ZooKeeper(ensemble, timeout_secs=cls.CONNECT_TIMEOUT_SECS, max_reconnects=cls.CONNECT_RETRIES, **kw)
def test_client_connect(): with ZookeeperClusterBootstrapper() as port: zk = ZooKeeper('localhost:%d' % port) assert zk.get_children('/') == ['zookeeper']