def make_endpoints(hostname, portmap, primary_port):
"""
Generate primary, additional endpoints from a portmap and primary_port.
primary_port must be a name in the portmap dictionary.
"""
# Do int check as stop-gap measure against incompatible downstream clients.
additional_endpoints = dict(
(name, Endpoint(hostname, port)) for (name, port) in portmap.items()
if isinstance(port, int))
# It's possible for the primary port to not have been allocated if this task
# is using autoregistration, so register with a port of 0.
return Endpoint(hostname, portmap.get(primary_port, 0)), additional_endpoints
def test_announcer_under_abnormal_circumstances():
mock_serverset = create_autospec(spec=ServerSet, instance=True)
mock_serverset.join = MagicMock()
mock_serverset.join.side_effect = [
KazooException('Whoops the ensemble is down!'),
'member0001',
]
mock_serverset.cancel = MagicMock()
endpoint = Endpoint('localhost', 12345)
clock = ThreadedClock(31337.0)
announcer = Announcer(mock_serverset,
endpoint,
clock=clock,
exception_wait=Amount(2, Time.SECONDS))
announcer.start()
try:
clock.tick(1.0)
assert announcer.disconnected_time() == 1.0
clock.tick(2.0)
assert announcer.disconnected_time() == 0.0, (
'Announcer should recover after an exception thrown internally.')
assert announcer._membership == 'member0001'
finally:
announcer.stop()
def test_make_empty_endpoints():
hostname = 'aurora.example.com'
portmap = {}
primary_port = 'http'
# test no bound 'http' port
primary, additional = make_endpoints(hostname, portmap, primary_port)
assert primary == Endpoint(hostname, 0)
assert additional == {}
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_announcer_on_expiration():
joined = threading.Event()
operations = []
def joined_side_effect(*args, **kw):
# 'global' does not work within python nested functions, so we cannot use a
# counter here, so instead we do append/len (see PEP-3104)
operations.append(1)
if len(operations) == 1 or len(operations) == 3:
joined.set()
return 'membership %d' % len(operations)
else:
raise KazooException('Failed to reconnect')
mock_serverset = create_autospec(spec=ServerSet, instance=True)
mock_serverset.join = MagicMock()
mock_serverset.join.side_effect = joined_side_effect
mock_serverset.cancel = MagicMock()
endpoint = Endpoint('localhost', 12345)
clock = ThreadedClock(31337.0)
announcer = Announcer(mock_serverset,
endpoint,
clock=clock,
exception_wait=Amount(2, Time.SECONDS))
announcer.start()
try:
joined.wait(timeout=1.0)
assert joined.is_set()
assert announcer._membership == 'membership 1'
assert announcer.disconnected_time() == 0.0
clock.tick(1.0)
assert announcer.disconnected_time() == 0.0
announcer.on_expiration() # expect exception
clock.tick(1.0)
assert announcer.disconnected_time() == 1.0, (
'Announcer should be disconnected on expiration.')
clock.tick(10.0)
assert announcer.disconnected_time() == 0.0, (
'Announcer should not advance disconnection time when connected.')
assert announcer._membership == 'membership 3'
finally:
announcer.stop()
def test_announcer_under_normal_circumstances():
joined = threading.Event()
def joined_side_effect(*args, **kw):
joined.set()
return 'membership foo'
mock_serverset = create_autospec(spec=ServerSet, instance=True)
mock_serverset.join = MagicMock()
mock_serverset.join.side_effect = joined_side_effect
mock_serverset.cancel = MagicMock()
endpoint = Endpoint('localhost', 12345)
clock = ThreadedClock(31337.0)
announcer = Announcer(mock_serverset, endpoint, clock=clock)
assert announcer.disconnected_time() == 0.0
clock.tick(1.0)
assert announcer.disconnected_time() == 1.0, (
'Announcer should advance disconnection time when not yet initially connected.'
)
announcer.start()
try:
joined.wait(timeout=1.0)
assert joined.is_set()
assert announcer.disconnected_time() == 0.0
clock.tick(1.0)
assert announcer.disconnected_time() == 0.0, (
'Announcer should not advance disconnection time when connected.')
assert announcer._membership == 'membership foo'
finally:
announcer.stop()
mock_serverset.cancel.assert_called_with('membership foo')
assert announcer.disconnected_time() == 0.0
clock.tick(1.0)
assert announcer.disconnected_time() == 0.0, (
'Announcer should not advance disconnection time when stopped.')
class ServerSetTestBase(object):
SERVICE_PATH = '/twitter/service/test'
INSTANCE1 = Endpoint(host='127.0.0.1', port=1234)
INSTANCE2 = Endpoint(host='127.0.0.1', port=1235)
ADDITIONAL1 = {'http': Endpoint(host='127.0.0.1', port=8080)}
ADDITIONAL2 = {'thrift': Endpoint(host='127.0.0.1', port=8081)}
@classmethod
def make_zk(cls, ensemble):
raise NotImplementedError
@classmethod
def session_id(cls, client):
raise NotImplementedError
def setUp(self):
self._server = ZookeeperServer()
def tearDown(self):
self._server.stop()
def test_client_iteration(self):
ss = ServerSet(self.make_zk(self._server.ensemble), self.SERVICE_PATH)
assert list(ss) == []
ss.join(self.INSTANCE1)
assert list(ss) == [ServiceInstance(self.INSTANCE1, member_id=0)]
ss.join(self.INSTANCE2)
assert list(ss) == [
ServiceInstance(self.INSTANCE1, member_id=0),
ServiceInstance(self.INSTANCE2, member_id=1)
]
def test_async_client_iteration(self):
ss1 = ServerSet(self.make_zk(self._server.ensemble), self.SERVICE_PATH)
ss2 = ServerSet(self.make_zk(self._server.ensemble), self.SERVICE_PATH)
ss1.join(self.INSTANCE1)
ss2.join(self.INSTANCE2)
assert list(ss1) == [
ServiceInstance(self.INSTANCE1, member_id=0),
ServiceInstance(self.INSTANCE2, member_id=1)
]
assert list(ss2) == [
ServiceInstance(self.INSTANCE1, member_id=0),
ServiceInstance(self.INSTANCE2, member_id=1)
]
def test_shard_id_registers(self):
ss1 = ServerSet(self.make_zk(self._server.ensemble), self.SERVICE_PATH)
ss2 = ServerSet(self.make_zk(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, member_id=0),
ServiceInstance(self.INSTANCE2, shard=1, member_id=1)
]
assert list(ss2) == [
ServiceInstance(self.INSTANCE1, shard=0, member_id=0),
ServiceInstance(self.INSTANCE2, shard=1, member_id=1)
]
def test_canceled_join_long_time(self):
zk = self.make_zk(self._server.ensemble)
zk.live.wait()
session_id = self.session_id(zk)
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_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(self.make_zk(self._server.ensemble),
self.SERVICE_PATH,
on_join=on_join,
on_leave=on_leave)
service2 = ServerSet(self.make_zk(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()
class TestServerSet(unittest.TestCase):
SERVICE_PATH = '/twitter/service/test'
INSTANCE1 = Endpoint(host='127.0.0.1', port=1234)
INSTANCE2 = Endpoint(host='127.0.0.1', port=1235)
ADDITIONAL1 = {'http': Endpoint(host='127.0.0.1', port=8080)}
ADDITIONAL2 = {'thrift': Endpoint(host='127.0.0.1', port=8081)}
def setUp(self):
self._server = ZookeeperServer()
def tearDown(self):
self._server.stop()
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 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_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()