class ZooWrap():
def __init__(self):
self.zk = KazooClient(hosts="%s,%s" % (
str(cc.conf['zookeeper']['host']),
str(cc.conf['zookeeper']['port'])))
self.zk.start()
self.root = cc.conf['zookeeper']['rootpath']
self.log = logging.getLogger('L.ZOOKEEPER')
self.zk.ensure_path('/%s/sleeping' % (self.root))
self.whoami = cc.conf['whoami']
def get_sleeping(self):
return self.zk.get_children('/%s/sleeping' % (self.root))
def sleep(self):
try:
self.zk.create('/%s/sleeping/%s' % (self.root, self.whoami))
self.log.info('Sleeping correctly')
except NodeExistsError:
self.log.error('Node already sleeping... seems weird')
def wake(self):
try:
self.zk.delete('/%s/sleeping/%s' % (self.root, self.whoami))
except NoNodeError:
self.log.error('Node was not sleeping... seems weird')
class ActorAddressBook(object):
def __init__(self, zk_hosts, timeout=60.0):
self.retry = KazooRetry(max_tries=10)
self.zk = KazooClient(hosts=zk_hosts, timeout=timeout)
self.zk.start()
def lookup(self, path):
return self.retry(self._lookup, path)
def _lookup(self, path):
actor_url, stat = self.zk.get(path)
return RemoteActor(actor_url.decode('utf-8'))
def register(self, path, actor_url):
return self.retry(self._register, path, actor_url)
def _register(self, path, actor_url):
self.zk.ensure_path(path)
self.zk.set(path, actor_url.encode('utf-8'))
def delete(self, path):
self.zk.delete(path, recursive=True)
def __del__(self):
self.zk.stop()
def cleanup(args):
now = dt.utcnow()
server = '{server}:{port}'.format(server=args.server, port=args.port)
logging.info('Connecting to {}'.format(server))
zk = KazooClient(hosts=server)
zk.start()
for path in args.zk_paths:
zk_path = '{}/{}'.format(args.zk_root_path, path)
nodes = zk.get_children(zk_path)
logging.info("Found {} nodes under {}".format(len(nodes), zk_path))
deleted = 0
for node in nodes:
node_path = '{}/{}'.format(zk_path, node)
data, stat = zk.get(node_path)
last_modified = dt.fromtimestamp(stat.mtime/1000.0)
if ((now - last_modified).days > args.age) or (args.inclusive and (now - last_modified).days >= args.age):
if not args.dry_run:
# Kazoo does not support recursive async deletes
if stat.children_count == 0:
res = zk.delete_async(node_path)
else:
zk.delete(node_path, recursive=True)
deleted += 1
logging.info("Deleted {} nodes".format(deleted))
zk.stop()
class ShellTestCase(unittest.TestCase):
""" base class for all tests """
def setUp(self):
"""
make sure that the prefix dir is empty
"""
self.tests_path = os.getenv("ZKSHELL_PREFIX_DIR", "/tests")
self.zk_host = os.getenv("ZKSHELL_ZK_HOST", "localhost:2181")
self.username = os.getenv("ZKSHELL_USER", "user")
self.password = os.getenv("ZKSHELL_PASSWD", "user")
self.digested_password = os.getenv("ZKSHELL_DIGESTED_PASSWD", "F46PeTVYeItL6aAyygIVQ9OaaeY=")
self.super_password = os.getenv("ZKSHELL_SUPER_PASSWD", "secret")
self.scheme = os.getenv("ZKSHELL_AUTH_SCHEME", "digest")
self.client = KazooClient(self.zk_host, 5)
self.client.start()
self.client.add_auth(self.scheme, self.auth_id)
if self.client.exists(self.tests_path):
self.client.delete(self.tests_path, recursive=True)
self.client.create(self.tests_path, str.encode(""))
self.output = StringIO()
self.shell = Shell([self.zk_host], 5, self.output, setup_readline=False, async=False)
# Create an empty test dir (needed for some tests)
self.temp_dir = tempfile.mkdtemp()
@property
def auth_id(self):
return "%s:%s" % (self.username, self.password)
@property
def auth_digest(self):
return "%s:%s" % (self.username, self.digested_password)
def tearDown(self):
self.output = None
self.shell = None
if os.path.isdir(self.temp_dir):
shutil.rmtree(self.temp_dir)
if self.client.exists(self.tests_path):
self.client.delete(self.tests_path, recursive=True)
self.client.stop()
###
# Helpers.
##
def create_compressed(self, path, value):
"""
ZK Shell doesn't support creating directly from a bytes array so we use a Kazoo client
to create a znode with zlib compressed content.
"""
compressed = zlib.compress(bytes(value, "utf-8") if PYTHON3 else value)
self.client.create(path, compressed, makepath=True)
class ZooKeeperTestMixin(object):
zk_hosts = None
_zk_hosts_internal = None
zk_base_path = None
proxy = None
def setup_zookeeper(self, base_path_prefix="/int_tests", use_proxy=False):
zk_hosts = os.environ.get("ZK_HOSTS")
if not zk_hosts:
raise unittest.SkipTest("export ZK_HOSTS env to run ZooKeeper integration tests")
if use_proxy:
hosts_list = zk_hosts.split(",")
if len(hosts_list) == 1:
self.proxy = SocatProxy(zk_hosts)
self.proxy.start()
self.zk_hosts = self.proxy.address
else:
proxies = [SocatProxy(host) for host in hosts_list]
self.proxy = MultiProxy(proxies)
self.proxy.start()
self.zk_hosts = ",".join(proxy.address for proxy in proxies)
self._zk_hosts_internal = zk_hosts
else:
self.zk_hosts = self._zk_hosts_internal = zk_hosts
self.zk_base_path = base_path_prefix + uuid.uuid4().hex
if os.environ.get("EPU_USE_GEVENT"):
from kazoo.handlers.gevent import SequentialGeventHandler
handler = SequentialGeventHandler()
self.use_gevent = True
else:
handler = None
self.use_gevent = False
self.kazoo = KazooClient(self._zk_hosts_internal + self.zk_base_path, handler=handler)
self.kazoo.start()
def teardown_zookeeper(self):
if self.kazoo:
try:
self.kazoo.delete("/", recursive=True)
self.kazoo.stop()
except Exception:
log.exception("Problem tearing down ZooKeeper")
if self.proxy:
self.proxy.stop()
cleanup_zookeeper = teardown_zookeeper
def _clean_up_zookeeper_autoscaling(context):
"""If max_instances was set for autoscaling, clean up zookeeper"""
client = KazooClient(hosts='%s/mesos-testcluster' % get_service_connection_string('zookeeper'), read_only=True)
client.start()
try:
client.delete('/autoscaling', recursive=True)
except NoNodeError:
pass
client.stop()
client.close()
def do_zookeeper_delete(addr, path):
print(path)
zk = KazooClient(addr)
zk.start()
zk.delete(path)
try:
do_zookeeper_read(addr, path)
except kazoo.exceptions.NoNodeError:
print('deleted')
class CockRoach(object):
def __init__(self, zkHost, stale_max_days=30, assume_yes=False, preview=False):
self.ConsumerGroups = []
self.zk_client = KazooClient(hosts=zkHost)
self.zk_client.start()
self.stale_max_days = stale_max_days
self.assume_yes = assume_yes
self.preview = preview
if self.zk_client.exists("/consumers"):
for cg_name in self.zk_client.get_children("/consumers"):
self.ConsumerGroups.append(ConsumerGroup(cg_name, \
self.zk_client))
def get_stale_cgroups(self, display):
"""
get_stale_cgroups returns ConsumerGroups
that were not used for stale_max_days
"""
ret = []
for consumergroup in self.ConsumerGroups:
delta = datetime.now() - consumergroup.last_seen().mtime
if delta.days > self.stale_max_days:
if display:
print "Stale: %s" % (consumergroup)
ret.append(consumergroup)
return ret
def delete_stale_cgroups(self):
""" Delete consumer groups that are considered stale"""
stale_cgroups = self.get_stale_cgroups(display=False)
for stale_cg in stale_cgroups:
print stale_cg
if self.assume_yes is False:
confirm = raw_input("Delete?")
else:
confirm = "Y"
if confirm == "Y":
self.delete_cgroup(stale_cg)
def delete_cgroup(self, consumergroup):
"""Deletes a consumer Group"""
print "Deleting %s" % (consumergroup.gid)
if self.preview is False:
self.zk_client.delete("/consumers/%s" % (consumergroup.gid), version=-1, recursive=True)
print "executed"
else:
print "pass"
def __str__(self):
ret = ""
for consumer in self.ConsumerGroups:
ret += "%s" % (consumer)
return ret
def delete_zk_data():
print "Deleting zk data..."
zk = KazooClient(hosts=config["zk"])
zk.start()
zk.delete("/kafka/brokers/topics", recursive=True)
zk.delete("/kafka/consumers", recursive=True)
zk.stop()
print "Deleted zk data."
def zk():
pytest.importorskip('kazoo')
pytest.mark.skipif(zookeeper_host is None,
reason='No ZOOKEEPER_HOST envar defined')
from kazoo.client import KazooClient
zk = KazooClient(hosts=zookeeper_host)
try:
zk.start()
yield zk
finally:
zk.delete('/satyr', recursive=True)
zk.stop()
def zookeeper_command(hosts, command, path):
try:
zk = KazooClient(hosts=hosts)
zk.start()
if command == 'get':
data, stat = zk.get(path)
return data.decode("utf-8")
elif command == 'delete':
zk.delete(path, recursive=True)
return 'Successfully deleted ' + path
else:
return False
zk.stop()
except:
return False
def test_create_makepath_incompatible_acls(self):
from kazoo.client import KazooClient
from kazoo.security import make_digest_acl_credential, CREATOR_ALL_ACL
credential = make_digest_acl_credential("username", "password")
alt_client = KazooClient(self.cluster[0].address + self.client.chroot,
max_retries=5, auth_data=[("digest", credential)])
alt_client.start()
alt_client.create("/1/2", b"val2", makepath=True, acl=CREATOR_ALL_ACL)
try:
self.assertRaises(NoAuthError, self.client.create, "/1/2/3/4/5",
b"val2", makepath=True)
finally:
alt_client.delete('/', recursive=True)
alt_client.stop()
def zk_command(self, command, path):
if self.client is None:
self._init_client()
zk_url = self.client.zk_url()
try:
zk = KazooClient(hosts=zk_url)
zk.start()
res = False
if command == "get":
data, stat = zk.get(path)
res = data.decode("utf-8")
elif command == "delete":
zk.delete(path, recursive=True)
res = "Successfully deleted " + path
zk.stop()
return res
except Exception as e:
self.vlog(e)
return False
class MasterElection(object):
def __init__(self):
self.zk = KazooClient(hosts=state.ARGS.zookeeper)
self.master_notified = False
self.my_node = None
self.zk.start() # Stop never called
self.zk.ensure_path(ELECTION_PATH)
def start_election(self, master_callback):
"""
Start the master election.
If this node is the master, the callback will be called once.
:params master_callback: Called if this node is the master
"""
self.callback = master_callback
self.my_node = self.zk.create(ELECTION_PATH + '/n_',
ephemeral=True, sequence=True)
self.zk.get_children(ELECTION_PATH, watch=self._election_update)
self._election_update()
def _election_update(self, data=None):
def worker():
try:
self.master_notified = True
self.callback()
except Exception as e:
self.master_notified = False
log.info("Failed to activate master, redoing election: %r", e)
self.zk.delete(self.my_node)
self.my_node = self.zk.create(ELECTION_PATH + '/n_',
ephemeral=True, sequence=True)
self._election_update()
if not self.master_notified and \
sorted(self.zk.get_children(ELECTION_PATH))[0] == \
os.path.basename(self.my_node):
t = threading.Thread(target=worker)
t.daemon = True
t.start()
class ZookeeperClient(object):
def __init__(self, zk_hosts):
self.zk = KazooClient(hosts=zk_hosts)
self.zk.start()
def put(self, entity, id, data):
path = '/{}/{}'.format(entity, id)
binary_value = json.dumps(data).encode('utf-8')
try:
self.zk.create(path, binary_value, makepath=True)
return True
except NodeExistsError:
self.zk.set(path, binary_value)
return False
def get(self, entity, id):
path = '/{}/{}'.format(entity, id)
try:
binary_data, _ = self.zk.get(path)
except NoNodeError:
return None
return json.loads(binary_data.decode('utf-8'))
def delete(self, entity, id):
path = '/{}/{}'.format(entity, id)
try:
self.zk.delete(path)
return True
except NoNodeError:
return False
def list(self, entity):
path = '/{}'.format(entity)
try:
children = self.zk.get_children(path)
for child in children:
value = self.get(entity, child)
yield (child, value)
except NoNodeError:
yield []
def test_init_servers(server_id, host = '127.0.0.1', port = 2181):
zookeeper = KazooClient('%s:%s' % (host, port,))
zookeeper.start()
try:
node = '/test/server_list'
if zookeeper.exists(node) is None:
zookeeper.create(node, json.dumps({'update_time' : time.time()}), makepath = True)
except kazoo.exceptions.NodeExistsError:
pass
try:
node = '/test/server_list/s%s' % server_id
zookeeper.delete(node)
zookeeper.create(node, json.dumps({
'update_time' : time.time(),
'server_name' : 's%s' % server_id,
'server_id' : server_id,
}), makepath = True)
except kazoo.exceptions.NodeExistsError:
pass
def main():
(options, args) = parse_options(sys.argv[1:])
data = options.znode_data_size * b"D"
s = KazooClient(options.server)
s.start()
if s.exists(options.root_znode):
children = s.get_children(options.root_znode)
print("delete old entries: %d" % len(children))
for child in children:
s.delete("%s/%s" % (options.root_znode, child))
else:
s.create(options.root_znode, "kazoo root znode")
evaluation(s, options.root_znode, data, options)
s.stop()
print("Performance test complete")
def zk_server(tmpdir):
zk_container_name = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(6))
# TODO(cmaloney): Add a python context manager for dockerized daemons
subprocess.check_call(['docker', 'run', '-d', '-p', '2181:2181', '-p',
'2888:2888', '-p', '3888:3888', '--name', zk_container_name, 'jplock/zookeeper'])
conn_retry_policy = KazooRetry(max_tries=-1, delay=0.1, max_delay=0.1)
cmd_retry_policy = KazooRetry(max_tries=3, delay=0.3, backoff=1, max_delay=1, ignore_expire=False)
zk = KazooClient(hosts=zk_hosts, connection_retry=conn_retry_policy, command_retry=cmd_retry_policy)
zk.start()
children = zk.get_children('/')
for child in children:
if child == 'zookeeper':
continue
zk.delete('/' + child, recursive=True)
yield zk
zk.stop()
zk.close()
subprocess.check_call(['docker', 'rm', '-f', zk_container_name])
class ZookClient(object):
def __init__(self, api_client=None):
logging.basicConfig()
# Create a client and start it
self.zk = KazooClient()
self.zk.start()
def create_accounts_path(self, name, **kwargs):
path = "/dso/" + name
self.zk.ensure_path(path)
self.zk.set(path, b"id: 7b4235ca-00fb-4dca-ad3e-8b6e3662631a\ngroupname: hr\ndescription: 人力资源")
def create_accountinfo_path(self, account_id, **kwargs):
self.zk.ensure_path("/app/someservice")
def create_path(self, path, **kwargs):
self.zk.ensure_path(path)
def get_data(self, path):
return self.zk.get(path)
def test_tran(self):
self.zk.delete("/app22")
self.zk.create("/app22", b"" + '{"12": "12"}')
tran = self.zk.transaction()
tran.delete("/app22")
tran.create("/app22", b"" + '{"22": "22"}')
tran.commit()
print "commit"
def stop(self):
# In the end, stop it
self.zk.stop()
class ClusterMonitor(threading.Thread):
"""periodically checks cluster member.
This class is delegated to change state between ACT clustered and ACT declustered."""
def __init__(self, zha):
threading.Thread.__init__(self)
self.zha = zha
self.should_run = True
self.zk = KazooClient(hosts=self.zha.config.get("connection_string","127.0.0.1:2181"), logger=logger)
self.zk.add_listener(self._zk_listener)
self.zk.start()
self.zroot = self.zha.config.get("cluster_znode","/zha-state")
self.znode = self.zroot + "/" + self.zha.config.get("id")
self._zk_register(first=True)
self.not_alone = None
def run(self):
while self.should_run:
time.sleep(self.zha.config.get("clustercheck_interval",3))
self.zha.recheck()
self._zk_register()
self.check_cluster()
self.trigger()
if self.zha.is_clustered:
self.zha.config.become_declustered()
self.zha.is_clustered = False
self.zk.delete(self.znode)
logger.info("cluster monitor thread stopped.")
def check_cluster(self):
try:
count = 0
chs = self.zk.get_children(self.zroot)
for ch in chs:
data, stats = self.zk.get(self.zroot+"/"+ch)
if data.strip()=="SBY:HEALTHY" and ch != self.zha.config.get("id"):
count += 1
if count != 0:
self.not_alone = time.time()
logger.debug("healthy sbys: %d"%(count,))
except Exception,e:
logger.warn("check cluster failed. Try next time.%s"%e)
class ZooDict(object):
def __init__(self, path_root=None):
self.path_root = path_root or default_path_root
def __enter__(self):
self.zk = KazooClient(hosts=zk_hosts)
self.zk.start(timeout=5)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.zk.stop()
def start(self):
return self.__enter__()
def stop(self):
return self.__exit__()
def get(self, key):
try:
r = self.zk.get(self.path_root + '/' + key)
if r and r[0]:
return json.loads(r[0])
except NoNodeError as no_node:
return None
def set(self, key, data_dict):
self.zk.ensure_path(self.path_root + '/' + key)
self.zk.set(self.path_root + '/' + key, json.dumps(data_dict))
def delete(self, key):
self.zk.delete(self.path_root + '/' + key)
def items(self):
return [ (child, self.get(child)) for child in self.zk.get_children(self.path_root) ]
class MicroService(object):
"""docstring for MicroService"""
def __init__(self, hosts, app):
from kazoo.client import KazooClient
self.zk = KazooClient(hosts=hosts)
self.app = "/" + app
def zkStart(self):
self.zk.start()
def zkStop(self):
self.zk.stop()
def registerService(self, service, url):
if self.zk.exists(self.app + "/" + service):
self.zk.delete(self.app + "/" + service, recursive=True)
self.zk.delete(self.app, recursive=True)
self.zk.ensure_path(self.app + "/" + service)
self.zk.set(self.app + "/" + service, bytes(url, encoding="utf-8"))
self.__watchServiceChanged(service)
return 'Service registed'
def unregisterService(self, service):
self.zk.delete(self.app + "/" + service, recursive=True)
return "Service deleted"
def findService(self, service):
data, stat = self.zk.get(self.app + "/" + service)
return data.decode("utf-8")
def __watchServiceChanged(self, service):
@self.zk.DataWatch(self.app + "/" + service)
def changed(data, stat):
print("Version: %s, data: %s" % (stat.version, data.decode("utf-8")))
next_temp=str(key)+" "+str(alphadict1[key])
clientSocket.send(next_temp)
clientSocket.close()
else :
pass
else :
connectionSocket.close()
else :
connectionSocket.close()
zk.delete("/alphanode")
else :
print("I might become the master!")
time.sleep(2)
if zk.exists("/masternode1") :
print("Someone else became the master so I can't! Going back to my regular duties")
os.system("gnome-terminal -e 'python Server2.py'")
quit()
else :
print 'I have become the master! I am ready to receive!I will carry out my regular duties and master duties!\n'
if zk.exists("/zlist") :
zk.set("/zlist","12000 /masternode 12000 /masternode1 65000 /numnode 45678 /alphanode1 45876 /specnode")
check_lis=zk.get("/zlist")
class ShellTestCase(unittest.TestCase):
""" base class for all tests """
@classmethod
def setUpClass(cls):
get_global_cluster().start()
def setUp(self):
"""
make sure that the prefix dir is empty
"""
self.tests_path = os.getenv("ZKSHELL_PREFIX_DIR", "/tests")
self.zk_hosts = ",".join(server.address
for server in get_global_cluster())
self.username = os.getenv("ZKSHELL_USER", "user")
self.password = os.getenv("ZKSHELL_PASSWD", "user")
self.digested_password = os.getenv("ZKSHELL_DIGESTED_PASSWD",
"F46PeTVYeItL6aAyygIVQ9OaaeY=")
self.super_password = os.getenv("ZKSHELL_SUPER_PASSWD", "secret")
self.scheme = os.getenv("ZKSHELL_AUTH_SCHEME", "digest")
self.client = KazooClient(self.zk_hosts, 5)
self.client.start()
self.client.add_auth(self.scheme, self.auth_id)
if self.client.exists(self.tests_path):
self.client.delete(self.tests_path, recursive=True)
self.client.create(self.tests_path, str.encode(""))
self.output = XStringIO()
self.shell = Shell([self.zk_hosts],
5,
self.output,
setup_readline=False,
async=False)
# Create an empty test dir (needed for some tests)
self.temp_dir = tempfile.mkdtemp()
@property
def auth_id(self):
return "%s:%s" % (self.username, self.password)
@property
def auth_digest(self):
return "%s:%s" % (self.username, self.digested_password)
def tearDown(self):
if self.output is not None:
self.output.close()
self.output = None
if self.shell is not None:
self.shell._disconnect()
self.shell = None
if os.path.isdir(self.temp_dir):
shutil.rmtree(self.temp_dir)
if self.client is not None:
if self.client.exists(self.tests_path):
self.client.delete(self.tests_path, recursive=True)
self.client.stop()
self.client.close()
self.client = None
###
# Helpers.
##
def create_compressed(self, path, value):
"""
ZK Shell doesn't support creating directly from a bytes array so we use a Kazoo client
to create a znode with zlib compressed content.
"""
compressed = zlib.compress(bytes(value, "utf-8") if PYTHON3 else value)
self.client.create(path, compressed, makepath=True)
# @Time : 2018/11/19 9:05
# @describe :
from kazoo.client import KazooClient
from kazoo.client import KazooState
import logging
import time
logging.basicConfig() # 解决建立连接异常因没有传入log日志对象而产生的报错
zk = KazooClient(hosts='39.107.88.145:2181', timeout=1) # 建立zk连接，超时时间设置为1s
@zk.add_listener # my_listener = zk.add_listener(my_listener)
def my_listener(state):
"""监听客户端连接状态
此函数看似没有显示调用，实际调用在装饰器中，也就是return时"""
if state == KazooState.LOST:
print 'lost'
elif state == KazooState.SUSPENDED:
print 'suspended'
else:
print 'connected'
zk.start()
# zk.set("/lws/test/1", b"some data test") # 更新给定节点数据，可以提供版本号，更新前需匹配
zk.delete("/lws/test/1", recursive=True) # 删除zNode节点，可选择递归删除，也可提供版本号删除
time.sleep(2)
zk.create("/lws/test/1", b"value one")
#!/usr/bin/env python
# coding=utf-8
# Jack Kang
from kazoo.client import KazooClient
zk = KazooClient("123.206.89.123:2181")
zk.start()
if zk.exists("/task/taobao") != None:
zk.delete('task/taobao', recursive=True)
if len(zk.get_children("/signal/taobao")) != 0:
zk.delete("/signal/taobao", recursive=True)
zk.create("/signal/taobao")
print "begin"
zk.create("/signal/taobao/start")
zk.stop()
class ZookeeperHandler:
def __init__(self, zk_hosts, prefix, host=None, port=None):
self.zk = KazooClient(hosts=zk_hosts)
self.prefix = prefix
self.is_service = host
self.children = None
self.root_path = self.prefix + "/{0}:{1}".format(host, port)
@self.zk.add_listener
def zk_listener(state):
# 这个方法也要放在start前面
if state == KazooState.LOST:
print("lost", state)
elif state == KazooState.CONNECTED:
# 检查是否是只读的客户端
if self.zk.client_state == KeeperState.CONNECTED_RO:
print("Read only mode!")
else:
print("Read/Write mode!")
elif state == KazooState.SUSPENDED:
print("SUSPENDED", state)
else:
print("other", state)
self.zk.start()
def register(self):
"""
服务注册
:return:
"""
if not self.is_service:
return
if not self.zk.exists(self.prefix):
# 注册这个目录，如果不存在的话就创建，可以连着创建几个目录（但是不能存储数据）
print("创建节点了", self.prefix)
self.zk.ensure_path(self.prefix)
if self.zk.exists(self.root_path):
self.zk.delete(self.root_path, recursive=True)
self.zk.create(self.root_path,
b'{"weight":1}',
ephemeral=True,
sequence=False)
# 监视子节点变化
@self.zk.ChildrenWatch(self.prefix)
def watch_children(children):
self.children = children
print("子节点变化了（需要更新服务列表了）: %s" % children)
# 实时监视数据的变化(可以做动态配置)
@self.zk.DataWatch(self.root_path)
def watch_node(data, stat):
print("节点下的数据变化了(这边可以写些处理逻辑)：", data, stat)
if stat:
print("Version: %s, data: %s" %
(stat.version, data.decode("utf-8")))
def discover(self):
"""
服务发现
:return:
"""
# 获取子节点
ret = self.zk.get_children(self.prefix)
return ret
class KazooCommandProxy():
def __init__(self, module):
self.module = module
self.zk = KazooClient(module.params['hosts'])
def absent(self):
return self._absent(self.module.params['name'])
def exists(self, znode):
return self.zk.exists(znode)
def list(self):
children = self.zk.get_children(self.module.params['name'])
return True, {'count': len(children), 'items': children, 'msg': 'Retrieved znodes in path.',
'znode': self.module.params['name']}
def present(self):
return self._present(self.module.params['name'], self.module.params['value'])
def get(self):
return self._get(self.module.params['name'])
def shutdown(self):
self.zk.stop()
self.zk.close()
def start(self):
self.zk.start()
def wait(self):
return self._wait(self.module.params['name'], self.module.params['timeout'])
def _absent(self, znode):
if self.exists(znode):
self.zk.delete(znode, recursive=self.module.params['recursive'])
return True, {'changed': True, 'msg': 'The znode was deleted.'}
else:
return True, {'changed': False, 'msg': 'The znode does not exist.'}
def _get(self, path):
if self.exists(path):
value, zstat = self.zk.get(path)
stat_dict = {}
for i in dir(zstat):
if not i.startswith('_'):
attr = getattr(zstat, i)
if isinstance(attr, (int, str)):
stat_dict[i] = attr
result = True, {'msg': 'The node was retrieved.', 'znode': path, 'value': value,
'stat': stat_dict}
else:
result = False, {'msg': 'The requested node does not exist.'}
return result
def _present(self, path, value):
if self.exists(path):
(current_value, zstat) = self.zk.get(path)
if value != current_value:
self.zk.set(path, to_bytes(value))
return True, {'changed': True, 'msg': 'Updated the znode value.', 'znode': path,
'value': value}
else:
return True, {'changed': False, 'msg': 'No changes were necessary.', 'znode': path, 'value': value}
else:
self.zk.create(path, to_bytes(value), makepath=True)
return True, {'changed': True, 'msg': 'Created a new znode.', 'znode': path, 'value': value}
def _wait(self, path, timeout, interval=5):
lim = time.time() + timeout
while time.time() < lim:
if self.exists(path):
return True, {'msg': 'The node appeared before the configured timeout.',
'znode': path, 'timeout': timeout}
else:
time.sleep(interval)
return False, {'msg': 'The node did not appear before the operation timed out.', 'timeout': timeout,
'znode': path}
class ZooKeeper(AbstractDCS):
def __init__(self, config):
super(ZooKeeper, self).__init__(config)
hosts = config.get('hosts', [])
if isinstance(hosts, list):
hosts = ','.join(hosts)
self._client = KazooClient(hosts, handler=PatroniSequentialThreadingHandler(config['retry_timeout']),
timeout=config['ttl'], connection_retry=KazooRetry(max_delay=1, max_tries=-1,
sleep_func=time.sleep), command_retry=KazooRetry(deadline=config['retry_timeout'],
max_delay=1, max_tries=-1, sleep_func=time.sleep))
self._client.add_listener(self.session_listener)
self._fetch_cluster = True
self._fetch_optime = True
self._orig_kazoo_connect = self._client._connection._connect
self._client._connection._connect = self._kazoo_connect
self._client.start()
def _kazoo_connect(self, *args):
"""Kazoo is using Ping's to determine health of connection to zookeeper. If there is no
response on Ping after Ping interval (1/2 from read_timeout) it will consider current
connection dead and try to connect to another node. Without this "magic" it was taking
up to 2/3 from session timeout (ttl) to figure out that connection was dead and we had
only small time for reconnect and retry.
This method is needed to return different value of read_timeout, which is not calculated
from negotiated session timeout but from value of `loop_wait`. And it is 2 sec smaller
than loop_wait, because we can spend up to 2 seconds when calling `touch_member()` and
`write_leader_optime()` methods, which also may hang..."""
ret = self._orig_kazoo_connect(*args)
return max(self.loop_wait - 2, 2)*1000, ret[1]
def session_listener(self, state):
if state in [KazooState.SUSPENDED, KazooState.LOST]:
self.cluster_watcher(None)
def optime_watcher(self, event):
self._fetch_optime = True
self.event.set()
def cluster_watcher(self, event):
self._fetch_cluster = True
self.optime_watcher(event)
def reload_config(self, config):
self.set_retry_timeout(config['retry_timeout'])
loop_wait = config['loop_wait']
loop_wait_changed = self._loop_wait != loop_wait
self._loop_wait = loop_wait
self._client.handler.set_connect_timeout(loop_wait)
# We need to reestablish connection to zookeeper if we want to change
# read_timeout (and Ping interval respectively), because read_timeout
# is calculated in `_kazoo_connect` method. If we are changing ttl at
# the same time, set_ttl method will reestablish connection and return
# `!True`, otherwise we will close existing connection and let kazoo
# open the new one.
if not self.set_ttl(int(config['ttl'] * 1000)) and loop_wait_changed:
self._client._connection._socket.close()
def set_ttl(self, ttl):
"""It is not possible to change ttl (session_timeout) in zookeeper without
destroying old session and creating the new one. This method returns `!True`
if session_timeout has been changed (`restart()` has been called)."""
if self._client._session_timeout != ttl:
self._client._session_timeout = ttl
self._client.restart()
return True
@property
def ttl(self):
return self._client._session_timeout
def set_retry_timeout(self, retry_timeout):
retry = self._client.retry if isinstance(self._client.retry, KazooRetry) else self._client._retry
retry.deadline = retry_timeout
def get_node(self, key, watch=None):
try:
ret = self._client.get(key, watch)
return (ret[0].decode('utf-8'), ret[1])
except NoNodeError:
return None
def get_leader_optime(self, leader):
watch = self.optime_watcher if not leader or leader.name != self._name else None
optime = self.get_node(self.leader_optime_path, watch)
self._fetch_optime = False
return optime and int(optime[0]) or 0
@staticmethod
def member(name, value, znode):
return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
def get_children(self, key, watch=None):
try:
return self._client.get_children(key, watch)
except NoNodeError:
return []
def load_members(self, sync_standby):
members = []
for member in self.get_children(self.members_path, self.cluster_watcher):
watch = member == sync_standby and self.cluster_watcher or None
data = self.get_node(self.members_path + member, watch)
if data is not None:
members.append(self.member(member, *data))
return members
def _inner_load_cluster(self):
self._fetch_cluster = False
self.event.clear()
nodes = set(self.get_children(self.client_path(''), self.cluster_watcher))
if not nodes:
self._fetch_cluster = True
# get initialize flag
initialize = (self.get_node(self.initialize_path) or [None])[0] if self._INITIALIZE in nodes else None
# get global dynamic configuration
config = self.get_node(self.config_path, watch=self.cluster_watcher) if self._CONFIG in nodes else None
config = config and ClusterConfig.from_node(config[1].version, config[0], config[1].mzxid)
# get timeline history
history = self.get_node(self.history_path, watch=self.cluster_watcher) if self._HISTORY in nodes else None
history = history and TimelineHistory.from_node(history[1].mzxid, history[0])
# get synchronization state
sync = self.get_node(self.sync_path, watch=self.cluster_watcher) if self._SYNC in nodes else None
sync = SyncState.from_node(sync and sync[1].version, sync and sync[0])
# get list of members
sync_standby = sync.leader == self._name and sync.sync_standby or None
members = self.load_members(sync_standby) if self._MEMBERS[:-1] in nodes else []
# get leader
leader = self.get_node(self.leader_path) if self._LEADER in nodes else None
if leader:
client_id = self._client.client_id
if not self._ctl and leader[0] == self._name and client_id is not None \
and client_id[0] != leader[1].ephemeralOwner:
logger.info('I am leader but not owner of the session. Removing leader node')
self._client.delete(self.leader_path)
leader = None
if leader:
member = Member(-1, leader[0], None, {})
member = ([m for m in members if m.name == leader[0]] or [member])[0]
leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
self._fetch_cluster = member.index == -1
# get last leader operation
last_leader_operation = self._OPTIME in nodes and self.get_leader_optime(leader)
# failover key
failover = self.get_node(self.failover_path, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
failover = failover and Failover.from_node(failover[1].version, failover[0])
return Cluster(initialize, config, leader, last_leader_operation, members, failover, sync, history)
def _load_cluster(self):
cluster = self.cluster
if self._fetch_cluster or cluster is None:
try:
cluster = self._client.retry(self._inner_load_cluster)
except Exception:
logger.exception('get_cluster')
self.cluster_watcher(None)
raise ZooKeeperError('ZooKeeper in not responding properly')
# Optime ZNode was updated or doesn't exist and we are not leader
elif (self._fetch_optime and not self._fetch_cluster or not cluster.last_leader_operation) and\
not (cluster.leader and cluster.leader.name == self._name):
try:
optime = self.get_leader_optime(cluster.leader)
cluster = Cluster(cluster.initialize, cluster.config, cluster.leader, optime,
cluster.members, cluster.failover, cluster.sync, cluster.history)
except Exception:
pass
return cluster
def _bypass_caches(self):
self._fetch_cluster = True
def _create(self, path, value, retry=False, ephemeral=False):
try:
if retry:
self._client.retry(self._client.create, path, value, makepath=True, ephemeral=ephemeral)
else:
self._client.create_async(path, value, makepath=True, ephemeral=ephemeral).get(timeout=1)
return True
except Exception:
logger.exception('Failed to create %s', path)
return False
def attempt_to_acquire_leader(self, permanent=False):
ret = self._create(self.leader_path, self._name.encode('utf-8'), retry=True, ephemeral=not permanent)
if not ret:
logger.info('Could not take out TTL lock')
return ret
def _set_or_create(self, key, value, index=None, retry=False, do_not_create_empty=False):
value = value.encode('utf-8')
try:
if retry:
self._client.retry(self._client.set, key, value, version=index or -1)
else:
self._client.set_async(key, value, version=index or -1).get(timeout=1)
return True
except NoNodeError:
if do_not_create_empty and not value:
return True
elif index is None:
return self._create(key, value, retry)
else:
return False
except Exception:
logger.exception('Failed to update %s', key)
return False
def set_failover_value(self, value, index=None):
return self._set_or_create(self.failover_path, value, index)
def set_config_value(self, value, index=None):
return self._set_or_create(self.config_path, value, index, retry=True)
def initialize(self, create_new=True, sysid=""):
sysid = sysid.encode('utf-8')
return self._create(self.initialize_path, sysid, retry=True) if create_new \
else self._client.retry(self._client.set, self.initialize_path, sysid)
def touch_member(self, data, permanent=False):
cluster = self.cluster
member = cluster and cluster.get_member(self._name, fallback_to_leader=False)
encoded_data = json.dumps(data, separators=(',', ':')).encode('utf-8')
if member and (self._client.client_id is not None and member.session != self._client.client_id[0] or
not (deep_compare(member.data.get('tags', {}), data.get('tags', {})) and
member.data.get('version') == data.get('version') and
member.data.get('checkpoint_after_promote') == data.get('checkpoint_after_promote'))):
try:
self._client.delete_async(self.member_path).get(timeout=1)
except NoNodeError:
pass
except Exception:
return False
member = None
if member:
if deep_compare(data, member.data):
return True
else:
try:
self._client.create_async(self.member_path, encoded_data, makepath=True,
ephemeral=not permanent).get(timeout=1)
return True
except Exception as e:
if not isinstance(e, NodeExistsError):
logger.exception('touch_member')
return False
try:
self._client.set_async(self.member_path, encoded_data).get(timeout=1)
return True
except Exception:
logger.exception('touch_member')
return False
def take_leader(self):
return self.attempt_to_acquire_leader()
def _write_leader_optime(self, last_operation):
return self._set_or_create(self.leader_optime_path, last_operation)
def _update_leader(self):
return True
def _delete_leader(self):
self._client.restart()
return True
def _cancel_initialization(self):
node = self.get_node(self.initialize_path)
if node:
self._client.delete(self.initialize_path, version=node[1].version)
def cancel_initialization(self):
try:
self._client.retry(self._cancel_initialization)
except Exception:
logger.exception("Unable to delete initialize key")
def delete_cluster(self):
try:
return self._client.retry(self._client.delete, self.client_path(''), recursive=True)
except NoNodeError:
return True
def set_history_value(self, value):
return self._set_or_create(self.history_path, value)
def set_sync_state_value(self, value, index=None):
return self._set_or_create(self.sync_path, value, index, retry=True, do_not_create_empty=True)
def delete_sync_state(self, index=None):
return self.set_sync_state_value("{}", index)
def watch(self, leader_index, timeout):
if super(ZooKeeper, self).watch(leader_index, timeout) and not self._fetch_optime:
self._fetch_cluster = True
return self._fetch_cluster
class ZkOperation(object):
def __init__(self, zk_hosts, zk_root):
self.zk = KazooClient(zk_hosts)
self.root = zk_root
self.tasks = set()
self.event = threading.Event()
def start(self):
if self.zk.exists:
self.zk.start()
self.zk.add_auth('digest', 'publish:publish')
if self.zk.connected:
self.zk.ensure_path(self.root)
def is_job_exist(self, job_id):
if job_id == '':
raise Exception('job_id is ""')
node = self.root + '/jobs/' + job_id
return self.zk.exists(node)
def check_task_status(self, path):
if path == '':
raise Exception('path is ""')
node = self.root + path
data, _ = self.zk.get(node)
return data['Status']
def _is_exist(self, node):
if self.zk.connected and self.zk.exists(node):
if self.zk.exists(node):
return True
else:
return False
else:
logger.error('zk not connected or node is exists')
return False
def _create_node(self, node, value=None):
if value is None:
value = ''
value = json.dumps(value)
if self.zk.connected and not self.zk.exists(node):
self.zk.create(node, makepath=True, value=value.encode())
return True
else:
logger.error('zk not connected or node is exists')
return False
def _update_node(self, node, value):
if self.zk.connected and self.zk.exists(node):
tx = self.zk.transaction()
tx.set_data(node, json.dumps(value).encode())
tx.commit()
return True
else:
logger.error(
'update node failed: zk not connected or node is not exists')
return False
def _get_node(self, node):
if self.zk.connected and self.zk.exists(node):
node_value, _ = self.zk.get(node)
if node_value:
return json.loads(node_value.decode())
else:
return {}
else:
logger.error('zk not connected or node is exists')
return False
def _delete_node(self, node):
if self.zk.connected:
if not self.zk.exists(node):
return True
else:
self.zk.delete(node, recursive=True)
return True
else:
logger.error('zk not connected')
return False
# is exist
def is_exist_signal(self, job_id):
node = '/{}/signal/{}'.format(self.root, job_id)
return self._is_exist(node)
# CREATE
def create_new_job(self, job_id, job_value=None):
if job_value is None:
job_value = ''
if job_id != '':
node = self.root + '/jobs/' + job_id
ret = self._create_node(node, job_value)
return ret
else:
logger.error('job_id is null')
return False
def create_new_target(self, job_id, target, target_value):
node = '/{}/jobs/{}/targets/{}'.format(self.root, job_id, target)
ret = self._create_node(node, target_value)
return ret
def create_new_task(self, job_id, target, task):
node = '/{}/jobs/{}/targets/{}/tasks/{}'.format(
self.root, job_id, target, task['task_id'])
ret = self._create_node(node, task)
return ret
def create_job_signal(self, job_id):
node = '/{}/signal/{}'.format(self.root, job_id)
ret = self._create_node(node, uuid.uuid4().hex)
return ret
# GET
def get_job_info(self, job_id):
job_node = '{}/jobs/{}'.format(self.root, job_id)
job_value, _ = self.zk.get(job_node)
job_info = json.loads(job_value.decode())
return job_info
def get_target_info(self, job_id, target):
target_node = '{}/jobs/{}/targets/{}'.format(self.root, job_id, target)
target_value, _ = self.zk.get(target_node)
target_info = json.loads(target_value.decode())
return target_info
def get_task_info(self, job_id, target, task_id):
task_node = '{}/jobs/{}/targets/{}/tasks/{}'.format(
self.root, job_id, target, task_id)
task_value, _ = self.zk.get(task_node)
task_info = json.loads(task_value.decode())
return task_info
# UPDATE
def update_job_status(self, job_id, task):
if job_id != '' and task is not None:
node = self.root + '/signal/' + job_id
else:
raise Exception('job_id is ""')
if self.zk.connected and self.is_job_exist(job_id):
tx = self.zk.transaction()
tx.set_data(node, task.encode())
tx.commit()
def handler_task(self, job_id, task_id, task_name, task_message, status):
# 为不必传回target, 遍历任务节点
if not self.is_job_exist(job_id):
logger.error("can not find this jobid: {}".format(job_id))
return False
job_node = "{}/jobs/{}/targets".format(self.root, job_id)
for target in self.zk.get_children(job_node):
target_node = "{}/{}/tasks".format(job_node, target)
for task in self.zk.get_children(target_node):
if task == task_id:
task_node = "{}/{}".format(target_node, task)
task_value, _ = self.zk.get(task_node)
new_task_value = json.loads(task_value.decode())
new_task_value['status'] = status
tx = self.zk.transaction()
tx.set_data(task_node, json.dumps(new_task_value).encode())
tx.commit()
task_value, _ = self.zk.get(task_node)
return True
logger.error("can not find this taskid: {} in {}".format(
task_id, job_id))
return False
def get_target_by_taskid(self, job_id, task_id):
if self.is_job_exist(job_id):
node = "{}/jobs/{}/targets".format(self.root, job_id)
for target in self.zk.get_children(node):
path = '{}/{}/tasks'.format(node, target)
for taskid in self.zk.get_children(path):
if taskid == task_id:
return target
return False
else:
logger.error("job is not exist: job_id={}".format(job_id))
def send_signal(self, job_id):
node = '{}/signal/{}'.format(self.root, job_id)
logger.info("send singal: {}".format(job_id))
tx = self.zk.transaction()
tx.set_data(node, uuid.uuid4().bytes)
tx.commit()
# DELETE
def delete_job(self, job_id):
node = "{}/jobs/{}".format(self.root, job_id)
logger.info("delete job: job_id={}".format(job_id))
self._delete_node(node)
def delete_signal(self, job_id):
node = '{}/signal/{}'.format(self.root, job_id)
logger.info("delete singal: {}".format(job_id))
self._delete_node(node)
def delete_target(self, job_id, target):
target_node = '{}/jobs/{}/targets/{}'.format(self.root, job_id, target)
logger.info("delete target: job_id={}, target={}".format(
job_id, target))
self._delete_node(target_node)
def delete_task(self, job_id, target, task_id):
task_node = '{}/jobs/{}/targets/{}/tasks/{}'.format(
self.root, job_id, target, task_id)
logger.info("delete task: job_id ={}, target={}, task_id={}".format(
job_id, target, task_id))
self._delete_node(task_node)
#################################
# CALLBACK
## exsit CALLBACK
def is_exist_callback(self, callback_node):
node = "{}/callback/{}".format(self.root, callback_node)
if self.zk.exists(node):
return True
else:
return False
## INIT CALLBACK
def init_callback_by_jobid(self, job_id):
node = "{}/callback/{}".format(self.root, job_id)
job_callback_value = {
"job_id": job_id,
"status": JobStatus.init.value,
"messages": ""
}
callback = {
"callback_level": CallbackLevel.job.value,
"callback_info": job_callback_value
}
ret = self._create_node(node, callback)
return ret
def init_callback_by_target(self, job_id, target):
node = "{}/callback/{}".format(self.root, job_id + "_" + target)
target_callback_value = {
"job_id": job_id,
"target": target,
"status": TargetStatus.init.value,
"messages": ""
}
callback = {
"callback_level": CallbackLevel.target.value,
"callback_info": target_callback_value
}
ret = self._create_node(node, callback)
return ret
def init_callback_by_taskid(self, job_id, target, task_id, task_name):
node = "{}/callback/{}".format(self.root, task_id)
taskid_callback_value = {
"job_id": job_id,
"target": target,
"task_name": task_name,
"status": JobStatus.init.value,
"messages": "",
}
callback = {
"callback_level": CallbackLevel.task.value,
"callback_info": taskid_callback_value
}
ret = self._create_node(node, callback)
return ret
## GET CALLBACK
def get_callback_info(self, callback):
node = "{}/callback/{}".format(self.root, callback)
if self.zk.exists(node):
node_value = self._get_node(node)
return node_value
else:
return False
## UPDATE CALLBACK
def update_callback_by_jobid(self, job_id, status, messages=None):
node = "{}/callback/{}".format(self.root, job_id)
if not self.zk.exists(node):
return False
node_value = self._get_node(node)
node_value["callback_info"]["status"] = status
if messages is not None:
node_value["callback_info"]["messages"] = messages
ret = self._update_node(node, node_value)
return ret
def update_callback_by_target(self, job_id, target, status, messages=None):
node = "{}/callback/{}".format(self.root, job_id + "_" + target)
if not self.zk.exists(node):
return False
node_value = self._get_node(node)
node_value["callback_info"]["status"] = status
if messages is not None:
node_value["callback_info"]["messages"] = messages
ret = self._update_node(node, node_value)
return ret
def update_callback_by_taskid(self, job_id, taskid, status, messages=None):
node = "{}/callback/{}".format(self.root, taskid)
if not self.zk.exists(node):
return False
node_value = self._get_node(node)
node_value["callback_info"]["status"] = status
if messages is not None:
node_value["callback_info"]["messages"] = messages
ret = self._update_node(node, node_value)
return ret
## DELETE CALLBACK
def delete_callback_node(self, callback):
node = "{}/callback/{}".format(self.root, callback)
ret = self._delete_node(node)
if ret:
logger.info(
"delete callback node success: callback={}".format(node))
else:
logger.error("delete callback node fail: callback={}".format(node))
return ret
default=False,
help='Delete items on /path recursively')
args = parser.parse_args()
zookeeper = {'server': '127.0.0.1', 'port': '2181'}
zk = KazooClient(hosts='%(server)s:%(port)s' % zookeeper)
zk.start()
if args.listing:
try:
answer = zk.get_children(args.item)
print(answer)
except kazoo.exceptions.NoNodeError:
print('ERROR: No such path!')
elif args.delete:
zk.delete(args.item, recursive=True)
print(args.item + ' deleted!')
else:
try:
answer = zk.get(args.item)
data = answer[0].decode('utf-8')
print(data)
except kazoo.exceptions.NoNodeError:
print('none')
zk.stop()
class ZK:
"""Opens a connection to a kafka zookeeper. "
"To be used in the 'with' statement."""
def __init__(self, cluster_config):
self.cluster_config = cluster_config
def __enter__(self):
kazooRetry = KazooRetry(max_tries=5, )
self.zk = KazooClient(
hosts=self.cluster_config.zookeeper,
read_only=True,
connection_retry=kazooRetry,
)
_log.debug(
"ZK: Creating new zookeeper connection: {zookeeper}".format(
zookeeper=self.cluster_config.zookeeper), )
self.zk.start()
return self
def __exit__(self, type, value, traceback):
self.zk.stop()
def get_children(self, path, watch=None):
"""Returns the children of the specified node."""
_log.debug("ZK: Getting children of {path}".format(path=path), )
return self.zk.get_children(path, watch)
def get(self, path, watch=None):
"""Returns the data of the specified node."""
_log.debug("ZK: Getting {path}".format(path=path), )
return self.zk.get(path, watch)
def set(self, path, value):
"""Sets and returns new data for the specified node."""
_log.debug("ZK: Setting {path} to {value}".format(path=path,
value=value))
return self.zk.set(path, value)
def get_json(self, path, watch=None):
"""Reads the data of the specified node and converts it to json."""
data, _ = self.get(path, watch)
return load_json(data) if data else None
def get_broker_metadata(self, broker_id):
try:
broker_json, _ = self.get(
"/brokers/ids/{b_id}".format(b_id=broker_id))
except NoNodeError:
_log.error("broker '{b_id}' not found.".format(b_id=broker_id), )
raise
return load_json(broker_json)
def get_brokers(self, names_only=False):
"""Get information on all the available brokers.
:rtype : dict of brokers
"""
try:
broker_ids = self.get_children("/brokers/ids")
except NoNodeError:
_log.info("cluster is empty.")
return {}
# Return broker-ids only
if names_only:
return {int(b_id): None for b_id in broker_ids}
return {
int(b_id): self.get_broker_metadata(b_id)
for b_id in broker_ids
}
def get_topic_config(self, topic):
"""Get configuration information for specified topic.
:rtype : dict of configuration
"""
try:
config_data = load_json(
self.get("/config/topics/{topic}".format(topic=topic))[0])
except NoNodeError as e:
# Kafka version before 0.8.1 does not have "/config/topics/<topic_name>" path in ZK and
# if the topic exists, return default dict instead of raising an Exception.
# Ref: https://cwiki.apache.org/confluence/display/KAFKA/Kafka+data+structures+in+Zookeeper.
topics = self.get_topics(topic_name=topic,
fetch_partition_state=False)
if len(topics) > 0:
_log.info(
"Configuration not available for topic {topic}.".format(
topic=topic))
config_data = {"config": {}}
else:
_log.error("topic {topic} not found.".format(topic=topic))
raise e
return config_data
def set_topic_config(self, topic, value, kafka_version=(
0,
10,
)):
"""Set configuration information for specified topic.
:topic : topic whose configuration needs to be changed
:value : config value with which the topic needs to be
updated with. This would be of the form key=value.
Example 'cleanup.policy=compact'
:kafka_version :tuple kafka version the brokers are running on.
Defaults to (0, 10, x). Kafka version 9 and kafka 10
support this feature.
"""
config_data = dump_json(value)
try:
# Change value
return_value = self.set(
"/config/topics/{topic}".format(topic=topic), config_data)
# Create change
version = kafka_version[1]
# this feature is supported in kafka 9 and kafka 10
assert version in (
9, 10), "Feature supported with kafka 9 and kafka 10"
if version == 9:
# https://github.com/apache/kafka/blob/0.9.0.1/
# core/src/main/scala/kafka/admin/AdminUtils.scala#L334
change_node = dump_json({
"version": 1,
"entity_type": "topics",
"entity_name": topic
})
else: # kafka 10
# https://github.com/apache/kafka/blob/0.10.2.1/
# core/src/main/scala/kafka/admin/AdminUtils.scala#L574
change_node = dump_json({
"version": 2,
"entity_path": "topics/" + topic,
})
self.create('/config/changes/config_change_',
change_node,
sequence=True)
except NoNodeError as e:
_log.error("topic {topic} not found.".format(topic=topic))
raise e
return return_value
def get_topics(
self,
topic_name=None,
names_only=False,
fetch_partition_state=True,
):
"""Get information on all the available topics.
Topic-data format with fetch_partition_state as False :-
topic_data = {
'version': 1,
'partitions': {
<p_id>: {
replicas: <broker-ids>
}
}
}
Topic-data format with fetch_partition_state as True:-
topic_data = {
'version': 1,
'ctime': <timestamp>,
'partitions': {
<p_id>:{
replicas: [<broker_id>, <broker_id>, ...],
isr: [<broker_id>, <broker_id>, ...],
controller_epoch: <val>,
leader_epoch: <val>,
version: 1,
leader: <broker-id>,
ctime: <timestamp>,
}
}
}
Note: By default we also fetch partition-state which results in
accessing the zookeeper twice. If just partition-replica information is
required fetch_partition_state should be set to False.
"""
try:
topic_ids = [topic_name] if topic_name else self.get_children(
"/brokers/topics", )
except NoNodeError:
_log.error("Cluster is empty.")
return {}
if names_only:
return topic_ids
topics_data = {}
for topic_id in topic_ids:
try:
topic_info = self.get(
"/brokers/topics/{id}".format(id=topic_id))
topic_data = load_json(topic_info[0])
topic_ctime = topic_info[1].ctime / 1000.0
topic_data['ctime'] = topic_ctime
except NoNodeError:
_log.info(
"topic '{topic}' not found.".format(topic=topic_id), )
return {}
# Prepare data for each partition
partitions_data = {}
for p_id, replicas in six.iteritems(topic_data['partitions']):
partitions_data[p_id] = {}
if fetch_partition_state:
# Fetch partition-state from zookeeper
partition_state = self._fetch_partition_state(
topic_id, p_id)
partitions_data[p_id] = load_json(partition_state[0])
partitions_data[p_id][
'ctime'] = partition_state[1].ctime / 1000.0
else:
# Fetch partition-info from zookeeper
partition_info = self._fetch_partition_info(topic_id, p_id)
partitions_data[p_id][
'ctime'] = partition_info.ctime / 1000.0
partitions_data[p_id]['replicas'] = replicas
topic_data['partitions'] = partitions_data
topics_data[topic_id] = topic_data
return topics_data
def get_consumer_groups(self, consumer_group_id=None, names_only=False):
"""Get information on all the available consumer-groups.
If names_only is False, only list of consumer-group ids are sent.
If names_only is True, Consumer group offset details are returned
for all consumer-groups or given consumer-group if given in dict
format as:-
{
'group-id':
{
'topic':
{
'partition': offset-value,
...
...
}
}
}
:rtype: dict of consumer-group offset details
"""
if consumer_group_id is None:
group_ids = self.get_children("/consumers")
else:
group_ids = [consumer_group_id]
# Return consumer-group-ids only
if names_only:
return {g_id: None for g_id in group_ids}
consumer_offsets = {}
for g_id in group_ids:
consumer_offsets[g_id] = self.get_group_offsets(g_id)
return consumer_offsets
def get_group_offsets(self, group, topic=None):
"""Fetch group offsets for given topic and partition otherwise all topics
and partitions otherwise.
{
'topic':
{
'partition': offset-value,
...
...
}
}
"""
group_offsets = {}
try:
all_topics = self.get_my_subscribed_topics(group)
except NoNodeError:
# No offset information of given consumer-group
_log.warning(
"No topics subscribed to consumer-group {group}.".format(
group=group, ), )
return group_offsets
if topic:
if topic in all_topics:
topics = [topic]
else:
_log.error(
"Topic {topic} not found in topic list {topics} for consumer"
"-group {consumer_group}.".format(
topic=topic,
topics=', '.join(topic for topic in all_topics),
consumer_group=group,
), )
return group_offsets
else:
topics = all_topics
for topic in topics:
group_offsets[topic] = {}
try:
partitions = self.get_my_subscribed_partitions(group, topic)
except NoNodeError:
_log.warning(
"No partition offsets found for topic {topic}. "
"Continuing to next one...".format(topic=topic), )
continue
# Fetch offsets for each partition
for partition in partitions:
path = "/consumers/{group_id}/offsets/{topic}/{partition}".format(
group_id=group,
topic=topic,
partition=partition,
)
try:
# Get current offset
offset_json, _ = self.get(path)
group_offsets[topic][partition] = load_json(offset_json)
except NoNodeError:
_log.error("Path {path} not found".format(path=path))
raise
return group_offsets
def _fetch_partition_state(self, topic_id, partition_id):
"""Fetch partition-state for given topic-partition."""
state_path = "/brokers/topics/{topic_id}/partitions/{p_id}/state"
try:
partition_state = self.get(
state_path.format(topic_id=topic_id, p_id=partition_id), )
return partition_state
except NoNodeError:
return {} # The partition has no data
def _fetch_partition_info(self, topic_id, partition_id):
"""Fetch partition info for given topic-partition."""
info_path = "/brokers/topics/{topic_id}/partitions/{p_id}"
try:
_, partition_info = self.get(
info_path.format(topic_id=topic_id, p_id=partition_id), )
return partition_info
except NoNodeError:
return {} # The partition has no data
def get_my_subscribed_topics(self, groupid):
"""Get the list of topics that a consumer is subscribed to
:param: groupid: The consumer group ID for the consumer
:returns list of kafka topics
:rtype: list
"""
path = "/consumers/{group_id}/offsets".format(group_id=groupid)
return self.get_children(path)
def get_my_subscribed_partitions(self, groupid, topic):
"""Get the list of partitions of a topic
that a consumer is subscribed to
:param: groupid: The consumer group ID for the consumer
:param: topic: The topic name
:returns list of partitions
:rtype: list
"""
path = "/consumers/{group_id}/offsets/{topic}".format(
group_id=groupid,
topic=topic,
)
return self.get_children(path)
def get_cluster_assignment(self):
"""Fetch the cluster layout in form of assignment from zookeeper"""
plan = self.get_cluster_plan()
assignment = {}
for elem in plan['partitions']:
assignment[(elem['topic'], elem['partition'])] = elem['replicas']
return assignment
def create(self,
path,
value='',
acl=None,
ephemeral=False,
sequence=False,
makepath=False):
"""Creates a Zookeeper node.
:param: path: The zookeeper node path
:param: value: Zookeeper node value
:param: acl: ACL list
:param: ephemeral: Boolean indicating where this node is tied to
this session.
:param: sequence: Boolean indicating whether path is suffixed
with a unique index.
:param: makepath: Whether the path should be created if it doesn't
exist.
"""
_log.debug("ZK: Creating node " + path)
return self.zk.create(path, value, acl, ephemeral, sequence, makepath)
def delete(self, path, recursive=False):
"""Deletes a Zookeeper node.
:param: path: The zookeeper node path
:param: recursive: Recursively delete node and all its children.
"""
_log.debug("ZK: Deleting node " + path)
return self.zk.delete(path, recursive=recursive)
def delete_topic_partitions(self, groupid, topic, partitions):
"""Delete the specified partitions within the topic that the consumer
is subscribed to.
:param: groupid: The consumer group ID for the consumer.
:param: topic: Kafka topic.
:param: partitions: List of partitions within the topic to be deleted.
:raises:
NoNodeError: if the consumer is not subscribed to the topic
ZookeeperError: if there is an error with Zookeeper
"""
for partition in partitions:
path = "/consumers/{groupid}/offsets/{topic}/{partition}".format(
groupid=groupid, topic=topic, partition=partition)
self.delete(path)
def delete_topic(self, groupid, topic):
path = "/consumers/{groupid}/offsets/{topic}".format(
groupid=groupid,
topic=topic,
)
self.delete(path, True)
def delete_group(self, groupid):
path = "/consumers/{groupid}".format(groupid=groupid, )
self.delete(path, True)
def execute_plan(self, plan, allow_rf_change=False):
"""Submit reassignment plan for execution."""
reassignment_path = '{admin}/{reassignment_node}'\
.format(admin=ADMIN_PATH, reassignment_node=REASSIGNMENT_NODE)
plan_json = dump_json(plan)
base_plan = self.get_cluster_plan()
if not validate_plan(plan, base_plan, allow_rf_change=allow_rf_change):
_log.error('Given plan is invalid. ABORTING reassignment...')
return False
# Send proposed-plan to zookeeper
try:
_log.info('Sending plan to Zookeeper...')
self.create(reassignment_path, plan_json, makepath=True)
_log.info(
'Re-assign partitions node in Zookeeper updated successfully '
'with {plan}'.format(plan=plan), )
return True
except NodeExistsError:
_log.warning('Previous plan in progress. Exiting..')
in_progress_plan = load_json(self.get(reassignment_path)[0])
in_progress_partitions = [
'{topic}-{p_id}'.format(
topic=p_data['topic'],
p_id=str(p_data['partition']),
) for p_data in in_progress_plan['partitions']
]
_log.warning(
'{count} partition(s) reassignment currently in progress:-'.
format(count=len(in_progress_partitions)), )
_log.warning(
'{partitions}. ABORTING reassignment...'.format(
partitions=', '.join(in_progress_partitions), ), )
return False
except Exception as e:
_log.error(
'Could not re-assign partitions {plan}. Error: {e}'.format(
plan=plan, e=e), )
return False
def get_cluster_plan(self):
"""Fetch cluster plan from zookeeper."""
_log.info('Fetching current cluster-topology from Zookeeper...')
cluster_layout = self.get_topics(fetch_partition_state=False)
# Re-format cluster-layout
partitions = [{
'topic': topic_id,
'partition': int(p_id),
'replicas': partitions_data['replicas']
} for topic_id, topic_info in six.iteritems(cluster_layout)
for p_id, partitions_data in six.iteritems(
topic_info['partitions'])]
return {'version': 1, 'partitions': partitions}
def get_pending_plan(self):
"""Read the currently running plan on reassign_partitions node."""
reassignment_path = '{admin}/{reassignment_node}'\
.format(admin=ADMIN_PATH, reassignment_node=REASSIGNMENT_NODE)
try:
result = self.get(reassignment_path)
return load_json(result[0])
except NoNodeError:
return {}
class ZkClient:
"""
Wrapper class over KazooClient. Provides utility methods for standalone failure tests to get details about
processor group state stored in zookeeper.
Instantiates a kazoo client to connect to zookeeper server at :param zookeeper_host::param zookeeper_port.
"""
def __init__(self, zookeeper_host, zookeeper_port, app_name, app_id):
self.kazoo_client = KazooClient(
hosts='{0}:{1}'.format(zookeeper_host, zookeeper_port))
self.zk_base_node = 'app-{0}-{1}/{2}-{3}-coordinationData'.format(
app_name, app_id, app_name, app_id)
def start(self):
"""
Establishes connection with the zookeeper server at self.host_name:self.port.
"""
self.kazoo_client.start()
def stop(self):
"""
Closes and releases the connection held with the zookeeper server.
"""
self.kazoo_client.stop()
def watch_job_model(self, watch_function):
self.kazoo_client.ensure_path(
'{0}/JobModelGeneration/jobModels/'.format(self.zk_base_node))
self.kazoo_client.get_children(
'{0}/JobModelGeneration/jobModels/'.format(self.zk_base_node),
watch=watch_function)
def get_latest_job_model(self):
"""
Reads and returns the latest JobModel from zookeeper.
"""
job_model_dict = {}
try:
childZkNodes = self.kazoo_client.get_children(
'{0}/JobModelGeneration/jobModels/'.format(self.zk_base_node))
if len(childZkNodes) > 0:
childZkNodes.sort()
childZkNodes.reverse()
job_model_generation_path = '{0}/JobModelGeneration/jobModels/{1}/'.format(
self.zk_base_node, childZkNodes[0])
job_model, _ = self.kazoo_client.get(job_model_generation_path)
"""
ZkClient java library stores the data in the following format in zookeeper:
class_name, data_length, actual_data
JobModel json manipulation: Delete all the characters before first occurrence of '{' in jobModel json string.
Normal json deserialization without the above custom string massaging fails. This will be removed after SAMZA-1749.
"""
first_curly_brace_index = job_model.find('{')
job_model = job_model[first_curly_brace_index:]
job_model_dict = json.loads(job_model)
logger.info(
"Recent JobModel in zookeeper: {0}".format(job_model_dict))
except:
logger.error(traceback.format_exc(sys.exc_info()))
return job_model_dict
def get_leader_processor_id(self):
"""
Determines the processorId of the current leader in a processors group.
Returns the processorId of the leader if leader exists.
Returns None otherwise.
"""
leader_processor_id = None
try:
processors_path = '{0}/processors'.format(self.zk_base_node)
childZkNodes = self.kazoo_client.get_children(processors_path)
childZkNodes.sort()
child_processor_path = '{0}/{1}'.format(processors_path,
childZkNodes[0])
processor_data, _ = self.kazoo_client.get(child_processor_path)
host, leader_processor_id = processor_data.split(" ")
except:
logger.error(traceback.format_exc(sys.exc_info()))
return leader_processor_id
def purge_all_nodes(self):
"""
Recursively delete all zookeeper nodes from the base node: self.zk_base_node.
"""
try:
self.kazoo_client.delete(path=self.zk_base_node,
version=-1,
recursive=True)
except:
logger.error(traceback.format_exc(sys.exc_info()))
def get_active_processors(self):
"""
Determines the processor ids that are active in zookeeper.
"""
processor_ids = []
try:
processors_path = '{0}/processors'.format(self.zk_base_node)
childZkNodes = self.kazoo_client.get_children(processors_path)
childZkNodes.sort()
for childZkNode in childZkNodes:
child_processor_path = '{0}/{1}'.format(
processors_path, childZkNode)
processor_data, _ = self.kazoo_client.get(child_processor_path)
host, processor_id = processor_data.split(" ")
processor_ids.append(processor_id)
except:
logger.error(traceback.format_exc(sys.exc_info()))
return processor_ids
# Handle being disconnected from Zookeeper
logger.error('session was suspended')
elif state == KazooState.CONNECTED:
logger.info('session is connected')
else:
# Handle being connected/reconnected to Zookeeper
logger.error('session was being connected/reconnected')
zk = KazooClient(hosts='127.0.0.1:2181')
zk.add_listener(my_listener)
zk.start()
id = int(time.time())
if zk.exists('/batches'):
zk.delete('/batches', recursive=True)
zk.ensure_path('/batches')
@zk.DataWatch('/batches/job_%s/node_1' % id)
def watch_node(data, stat):
try:
print("Version: %s, data: %s" % (stat.version, data.decode("utf-8")))
except Exception as e:
print(e)
@zk.ChildrenWatch('/batches')
def watch_children(children):
try:
class zookeeper:
def __init__(self, address):
self.address = address
self.zk = KazooClient(address)
self.zk.start()
self.arcus_cache_map = {}
self.arcus_node_map = {}
self.force = False
self.meta = ('', None)
self.meta_mtime = None
def __repr__(self):
repr = '[ZooKeeper: %s] %s, %s' % (self.address, self.meta[0], str(self.meta[1]))
for code, cache in self.arcus_cache_map.items():
repr = '%s\n\n%s' % (repr, cache)
return repr
def set_force(self):
self.force = True
def zk_read(self, path):
data, stat = self.zk.get(path)
children = self.zk.get_children(path)
return data, stat, children
def zk_children(self, path, watch=None):
if watch != None:
return self.zk.get_children(path, watch = watch)
else:
return self.zk.get_children(path)
def zk_children_if_exists(self, path, watch=None):
if self.zk_exists(path) == False:
return []
return self.zk_children(path, watch)
def zk_exists(self, path):
if self.zk.exists(path) == None:
return False
return True
def zk_create(self, path, value):
try:
self.zk.create(path, bytes(value, 'utf-8'))
except NodeExistsError:
if self.force == False:
raise NodeExistsError
def zk_delete(self, path):
try:
self.zk.delete(path)
except NoNodeError:
if self.force == False:
raise NoNodeError
def zk_delete_tree(self, path):
try:
self.zk.delete(path, recursive=True)
except NoNodeError:
if self.force == False:
raise NoNodeError
def zk_update(self, path, value):
try:
self.zk.set(path, bytes(value, 'utf-8'))
except NoNodeError:
if self.force == False:
raise NoNodeError
def get_arcus_cache_list(self):
children = self.zk_children_if_exists('/arcus/cache_list/')
children += self.zk_children_if_exists('/arcus_repl/cache_list/')
return children
def get_arcus_node_of_code(self, code, server):
# repl case
children = self.zk_children_if_exists('/arcus_repl/cache_list/' + code)
children += self.zk_children_if_exists('/arcus/cache_list/' + code)
ret = []
for child in children:
tmp = child.split('^', 2) # remove repl info
if len(tmp) == 3:
child = tmp[2]
addr, name = child.split('-', 1)
ip, port = addr.split(':', 1)
if server != '' and (server != ip and server != name):
continue # skip this
node = arcus_node(ip, port)
node.name = name
ret.append(node)
return ret
def get_arcus_node_of_server(self, addr):
ip = socket.gethostbyname(addr)
children = self.zk_children_if_exists('/arcus_repl/cache_server_mapping/')
children += self.zk_children_if_exists('/arcus/cache_server_mapping/')
ret = []
for child in children:
l = len(ip)
if child[:l] == ip:
code = self.zk_children_if_exists('/arcus_repl/cache_server_mapping/' + child)
if len(code) == 0:
code = self.zk_children_if_exists('/arcus/cache_server_mapping/' + child)
code = code[0]
tmp = code.split('^') # remove repl info
code = tmp[0]
try:
ip, port = child.split(':')
except ValueError:
print('No port defined in cache_server_mapping: %s' % child)
continue
node = arcus_node(ip, port)
node.code = code
ret.append(node)
return ret
def _get_arcus_node(self, child, results):
code = self.zk_children_if_exists('/arcus_repl/cache_server_mapping/' + child)
if len(code) == 0:
code = self.zk_children_if_exists('/arcus/cache_server_mapping/' + child)
if len(code) == 0:
print('no childrens in cache_server_mapping error: %s' % child)
print(code)
return
code = code[0]
tmp = code.split('^') # remove repl info
code = tmp[0]
try:
ip, port = child.split(':')
except ValueError:
print('No port defined in cache_server_mapping: %s' % child)
ip = child
port = '0'
node = arcus_node(ip, port)
node.code = code
results.append(node)
def get_arcus_node_all(self):
children = self.zk_children_if_exists('/arcus_repl/cache_server_mapping/')
children += self.zk_children_if_exists('/arcus/cache_server_mapping/')
ret = []
threads = []
#print(children)
for child in children:
th = threading.Thread(target = self._get_arcus_node, args = (child, ret))
th.start()
threads.append(th)
for th in threads:
th.join()
return ret
def _get_arcus_meta(self, child, results):
data, stat, children = self.zk_read('/arcus/meta/' + child)
results[child] = [data.decode('utf-8'), stat]
def get_arcus_meta_all(self):
if self.zk.exists('/arcus/meta') == None:
self.zk.create('/arcus/meta', b'arcus meta info')
children = self.zk.get_children('/arcus/meta')
print('# children')
print(children)
threads = []
ret = {}
#print(children)
for child in children:
th = threading.Thread(target = self._get_arcus_meta, args = (child, ret))
th.start()
threads.append(th)
for th in threads:
th.join()
return ret
def _match_code_and_nodes(self, code, cache, meta):
#repl case
children = self.zk_children_if_exists('/arcus_repl/cache_list/' + code)
children += self.zk_children_if_exists('/arcus/cache_list/' + code)
for child in children:
tmp = child.split('^', 2) # remove repl info
if len(tmp) == 3:
child = tmp[2]
addr, name = child.split('-')
try:
node = self.arcus_node_map[addr]
except KeyError:
print('%s of %s is not defined in cache_server_mapping' % (addr, code))
ip, port = addr.split(':')
node = arcus_node(ip, port)
node.noport = True
node.active = True
cache.active_node.append(node)
for node in cache.node:
if node.active == False:
cache.dead_node.append(node)
if code in meta:
cache.meta = meta[code]
def load_all(self):
codes = self.get_arcus_cache_list()
for code in codes:
cache = arcus_cache(self.address, code)
self.arcus_cache_map[code] = cache
print('# get_arcus_node_all()')
nodes = self.get_arcus_node_all()
print('# done')
for node in nodes:
self.arcus_node_map[node.ip + ":" + node.port] = node
self.arcus_cache_map[node.code].node.append(node)
# meta info
print('# get_arcus_meta_all()')
meta = self.get_arcus_meta_all()
print('# done')
print('# match code & nodes')
threads = []
for code, cache in self.arcus_cache_map.items():
th = threading.Thread(target = self._match_code_and_nodes, args = (code, cache, meta))
th.start()
threads.append(th)
for th in threads:
th.join()
print('#done')
if 'zookeeper' in meta:
self.meta = meta['zookeeper']
def _callback(self, event):
child_list = self.zk.get_children(event.path)
cloud = os.path.basename(event.path)
cache = self.arcus_cache_map[cloud]
event_list = { 'created':[], 'deleted':[] }
current = {}
print('##### active node')
print(cache.active_node)
children = []
for child in child_list:
addr = child.split('-')[0]
children.append(addr)
print('#### children')
print(children)
for node in cache.active_node:
current[node.ip + ':' + node.port] = True
print('##### current')
print(current)
for node in cache.active_node:
addr = node.ip + ':' + node.port
if addr not in children:
event_list['deleted'].append(addr)
cache.active_node.remove(node)
for child in children:
if child not in current:
event_list['created'].append(child)
ip, port = child.split(':')
node = arcus_node(ip, port)
cache.active_node.append(node)
print('####### result')
print(cache.active_node)
self.callback(event, event_list)
children = self.zk.get_children(event.path, watch = self._callback)
def watch(self, callback):
self.callback = callback
for code, cache in self.arcus_cache_map.items():
children = self.zk_children_if_exists('/arcus/cache_list/' + code, watch=self._callback)
children += self.zk_children_if_exists('/arcus_repl/cache_list/' + code, watch=self._callback)
class BalancedConsumer(object):
"""
A self-balancing consumer for Kafka that uses ZooKeeper to communicate
with other balancing consumers.
Maintains a single instance of SimpleConsumer, periodically using the
consumer rebalancing algorithm to reassign partitions to this
SimpleConsumer.
"""
def __init__(self,
topic,
cluster,
consumer_group,
fetch_message_max_bytes=1024 * 1024,
num_consumer_fetchers=1,
auto_commit_enable=False,
auto_commit_interval_ms=60 * 1000,
queued_max_messages=2000,
fetch_min_bytes=1,
fetch_wait_max_ms=100,
offsets_channel_backoff_ms=1000,
offsets_commit_max_retries=5,
auto_offset_reset=OffsetType.EARLIEST,
consumer_timeout_ms=-1,
rebalance_max_retries=5,
rebalance_backoff_ms=2 * 1000,
zookeeper_connection_timeout_ms=6 * 1000,
zookeeper_connect='127.0.0.1:2181',
zookeeper=None,
auto_start=True,
reset_offset_on_start=False,
post_rebalance_callback=None,
use_rdkafka=False):
"""Create a BalancedConsumer instance
:param topic: The topic this consumer should consume
:type topic: :class:`pykafka.topic.Topic`
:param cluster: The cluster to which this consumer should connect
:type cluster: :class:`pykafka.cluster.Cluster`
:param consumer_group: The name of the consumer group this consumer
should join.
:type consumer_group: bytes
:param fetch_message_max_bytes: The number of bytes of messages to
attempt to fetch with each fetch request
:type fetch_message_max_bytes: int
:param num_consumer_fetchers: The number of workers used to make
FetchRequests
:type num_consumer_fetchers: int
:param auto_commit_enable: If true, periodically commit to kafka the
offset of messages already fetched by this consumer. This also
requires that `consumer_group` is not `None`.
:type auto_commit_enable: bool
:param auto_commit_interval_ms: The frequency (in milliseconds) at which
the consumer's offsets are committed to kafka. This setting is
ignored if `auto_commit_enable` is `False`.
:type auto_commit_interval_ms: int
:param queued_max_messages: The maximum number of messages buffered for
consumption in the internal
:class:`pykafka.simpleconsumer.SimpleConsumer`
:type queued_max_messages: int
:param fetch_min_bytes: The minimum amount of data (in bytes) that the
server should return for a fetch request. If insufficient data is
available, the request will block until sufficient data is available.
:type fetch_min_bytes: int
:param fetch_wait_max_ms: The maximum amount of time (in milliseconds)
that the server will block before answering a fetch request if
there isn't sufficient data to immediately satisfy `fetch_min_bytes`.
:type fetch_wait_max_ms: int
:param offsets_channel_backoff_ms: Backoff time to retry failed offset
commits and fetches.
:type offsets_channel_backoff_ms: int
:param offsets_commit_max_retries: The number of times the offset commit
worker should retry before raising an error.
:type offsets_commit_max_retries: int
:param auto_offset_reset: What to do if an offset is out of range. This
setting indicates how to reset the consumer's internal offset
counter when an `OffsetOutOfRangeError` is encountered.
:type auto_offset_reset: :class:`pykafka.common.OffsetType`
:param consumer_timeout_ms: Amount of time (in milliseconds) the
consumer may spend without messages available for consumption
before returning None.
:type consumer_timeout_ms: int
:param rebalance_max_retries: The number of times the rebalance should
retry before raising an error.
:type rebalance_max_retries: int
:param rebalance_backoff_ms: Backoff time (in milliseconds) between
retries during rebalance.
:type rebalance_backoff_ms: int
:param zookeeper_connection_timeout_ms: The maximum time (in
milliseconds) that the consumer waits while establishing a
connection to zookeeper.
:type zookeeper_connection_timeout_ms: int
:param zookeeper_connect: Comma-separated (ip1:port1,ip2:port2) strings
indicating the zookeeper nodes to which to connect.
:type zookeeper_connect: str
:param zookeeper: A KazooClient connected to a Zookeeper instance.
If provided, `zookeeper_connect` is ignored.
:type zookeeper: :class:`kazoo.client.KazooClient`
:param auto_start: Whether the consumer should begin communicating
with zookeeper after __init__ is complete. If false, communication
can be started with `start()`.
:type auto_start: bool
:param reset_offset_on_start: Whether the consumer should reset its
internal offset counter to `self._auto_offset_reset` and commit that
offset immediately upon starting up
:type reset_offset_on_start: bool
:param post_rebalance_callback: A function to be called when a rebalance is
in progress. This function should accept three arguments: the
:class:`pykafka.balancedconsumer.BalancedConsumer` instance that just
completed its rebalance, a dict of partitions that it owned before the
rebalance, and a dict of partitions it owns after the rebalance. These dicts
map partition ids to the most recently known offsets for those partitions.
This function can optionally return a dictionary mapping partition ids to
offsets. If it does, the consumer will reset its offsets to the supplied
values before continuing consumption.
Note that the BalancedConsumer is in a poorly defined state at
the time this callback runs, so that accessing its properties
(such as `held_offsets` or `partitions`) might yield confusing
results. Instead, the callback should really rely on the
provided partition-id dicts, which are well-defined.
:type post_rebalance_callback: function
:param use_rdkafka: Use librdkafka-backed consumer if available
:type use_rdkafka: bool
"""
self._cluster = cluster
self._consumer_group = consumer_group
self._topic = topic
self._auto_commit_enable = auto_commit_enable
self._auto_commit_interval_ms = auto_commit_interval_ms
self._fetch_message_max_bytes = fetch_message_max_bytes
self._fetch_min_bytes = fetch_min_bytes
self._rebalance_max_retries = rebalance_max_retries
self._num_consumer_fetchers = num_consumer_fetchers
self._queued_max_messages = queued_max_messages
self._fetch_wait_max_ms = fetch_wait_max_ms
self._rebalance_backoff_ms = rebalance_backoff_ms
self._consumer_timeout_ms = consumer_timeout_ms
self._offsets_channel_backoff_ms = offsets_channel_backoff_ms
self._offsets_commit_max_retries = offsets_commit_max_retries
self._auto_offset_reset = auto_offset_reset
self._zookeeper_connect = zookeeper_connect
self._zookeeper_connection_timeout_ms = zookeeper_connection_timeout_ms
self._reset_offset_on_start = reset_offset_on_start
self._post_rebalance_callback = post_rebalance_callback
self._running = False
self._worker_exception = None
self._worker_trace_logged = False
if not rdkafka and use_rdkafka:
raise ImportError("use_rdkafka requires rdkafka to be installed")
self._use_rdkafka = rdkafka and use_rdkafka
self._rebalancing_lock = cluster.handler.Lock()
self._consumer = None
self._consumer_id = "{hostname}:{uuid}".format(
hostname=socket.gethostname(),
uuid=uuid4()
)
self._setting_watches = True
self._topic_path = '/consumers/{group}/owners/{topic}'.format(
group=self._consumer_group,
topic=self._topic.name)
self._consumer_id_path = '/consumers/{group}/ids'.format(
group=self._consumer_group)
self._zookeeper = None
self._owns_zookeeper = zookeeper is None
if zookeeper is not None:
self._zookeeper = zookeeper
if auto_start is True:
self.start()
def __del__(self):
log.debug("Finalising {}".format(self))
self.stop()
def __repr__(self):
return "<{module}.{name} at {id_} (consumer_group={group})>".format(
module=self.__class__.__module__,
name=self.__class__.__name__,
id_=hex(id(self)),
group=self._consumer_group
)
def _raise_worker_exceptions(self):
"""Raises exceptions encountered on worker threads"""
if self._worker_exception is not None:
_, ex, tb = self._worker_exception
if not self._worker_trace_logged:
self._worker_trace_logged = True
log.error("Exception encountered in worker thread:\n%s",
"".join(traceback.format_tb(tb)))
raise ex
def _setup_checker_worker(self):
"""Start the zookeeper partition checker thread"""
self = weakref.proxy(self)
def checker():
while True:
try:
if not self._running:
break
time.sleep(120)
if not self._check_held_partitions():
self._rebalance()
except Exception as e:
if not isinstance(e, ReferenceError):
# surface all exceptions to the main thread
self._worker_exception = sys.exc_info()
break
log.debug("Checker thread exiting")
log.debug("Starting checker thread")
return self._cluster.handler.spawn(checker)
@property
def partitions(self):
return self._consumer.partitions if self._consumer else dict()
@property
def _partitions(self):
"""Convenient shorthand for set of partitions internally held"""
return set(
[] if self.partitions is None else itervalues(self.partitions))
@property
def held_offsets(self):
"""Return a map from partition id to held offset for each partition"""
if not self._consumer:
return None
return self._consumer.held_offsets
def start(self):
"""Open connections and join a cluster."""
try:
if self._zookeeper is None:
self._setup_zookeeper(self._zookeeper_connect,
self._zookeeper_connection_timeout_ms)
self._zookeeper.ensure_path(self._topic_path)
self._add_self()
self._running = True
self._set_watches()
self._rebalance()
self._setup_checker_worker()
except Exception:
log.error("Stopping consumer in response to error")
self.stop()
def stop(self):
"""Close the zookeeper connection and stop consuming.
This method should be called as part of a graceful shutdown process.
"""
log.debug("Stopping {}".format(self))
with self._rebalancing_lock:
# We acquire the lock in order to prevent a race condition where a
# rebalance that is already underway might re-register the zk
# nodes that we remove here
self._running = False
if self._consumer is not None:
self._consumer.stop()
if self._owns_zookeeper:
# NB this should always come last, so we do not hand over control
# of our partitions until consumption has really been halted
self._zookeeper.stop()
else:
self._remove_partitions(self._get_held_partitions())
try:
self._zookeeper.delete(self._path_self)
except NoNodeException:
pass
# additionally we'd want to remove watches here, but there are no
# facilities for that in ChildrenWatch - as a workaround we check
# self._running in the watcher callbacks (see further down)
def _setup_zookeeper(self, zookeeper_connect, timeout):
"""Open a connection to a ZooKeeper host.
:param zookeeper_connect: The 'ip:port' address of the zookeeper node to
which to connect.
:type zookeeper_connect: str
:param timeout: Connection timeout (in milliseconds)
:type timeout: int
"""
self._zookeeper = KazooClient(zookeeper_connect, timeout=timeout / 1000)
self._zookeeper.start()
def _setup_internal_consumer(self, partitions=None, start=True):
"""Instantiate an internal SimpleConsumer instance"""
self._consumer = self._get_internal_consumer(partitions=partitions, start=start)
def _get_internal_consumer(self, partitions=None, start=True):
"""Instantiate a SimpleConsumer for internal use.
If there is already a SimpleConsumer instance held by this object,
disable its workers and mark it for garbage collection before
creating a new one.
"""
if partitions is None:
partitions = []
reset_offset_on_start = self._reset_offset_on_start
if self._consumer is not None:
self._consumer.stop()
# only use this setting for the first call to
# _get_internal_consumer. subsequent calls should not
# reset the offsets, since they can happen at any time
reset_offset_on_start = False
Cls = (rdkafka.RdKafkaSimpleConsumer
if self._use_rdkafka else SimpleConsumer)
return Cls(
self._topic,
self._cluster,
consumer_group=self._consumer_group,
partitions=partitions,
auto_commit_enable=self._auto_commit_enable,
auto_commit_interval_ms=self._auto_commit_interval_ms,
fetch_message_max_bytes=self._fetch_message_max_bytes,
fetch_min_bytes=self._fetch_min_bytes,
num_consumer_fetchers=self._num_consumer_fetchers,
queued_max_messages=self._queued_max_messages,
fetch_wait_max_ms=self._fetch_wait_max_ms,
consumer_timeout_ms=self._consumer_timeout_ms,
offsets_channel_backoff_ms=self._offsets_channel_backoff_ms,
offsets_commit_max_retries=self._offsets_commit_max_retries,
auto_offset_reset=self._auto_offset_reset,
reset_offset_on_start=reset_offset_on_start,
auto_start=start
)
def _decide_partitions(self, participants):
"""Decide which partitions belong to this consumer.
Uses the consumer rebalancing algorithm described here
http://kafka.apache.org/documentation.html
It is very important that the participants array is sorted,
since this algorithm runs on each consumer and indexes into the same
array. The same array index operation must return the same
result on each consumer.
:param participants: Sorted list of ids of all other consumers in this
consumer group.
:type participants: Iterable of `bytes`
"""
# Freeze and sort partitions so we always have the same results
p_to_str = lambda p: '-'.join([str(p.topic.name), str(p.leader.id), str(p.id)])
all_parts = self._topic.partitions.values()
all_parts = sorted(all_parts, key=p_to_str)
# get start point, # of partitions, and remainder
participants = sorted(participants) # just make sure it's sorted.
idx = participants.index(self._consumer_id)
parts_per_consumer = len(all_parts) // len(participants)
remainder_ppc = len(all_parts) % len(participants)
start = parts_per_consumer * idx + min(idx, remainder_ppc)
num_parts = parts_per_consumer + (0 if (idx + 1 > remainder_ppc) else 1)
# assign partitions from i*N to (i+1)*N - 1 to consumer Ci
new_partitions = itertools.islice(all_parts, start, start + num_parts)
new_partitions = set(new_partitions)
log.info('Balancing %i participants for %i partitions.\nOwning %i partitions.',
len(participants), len(all_parts), len(new_partitions))
log.debug('My partitions: %s', [p_to_str(p) for p in new_partitions])
return new_partitions
def _get_participants(self):
"""Use zookeeper to get the other consumers of this topic.
:return: A sorted list of the ids of the other consumers of this
consumer's topic
"""
try:
consumer_ids = self._zookeeper.get_children(self._consumer_id_path)
except NoNodeException:
log.debug("Consumer group doesn't exist. "
"No participants to find")
return []
participants = []
for id_ in consumer_ids:
try:
topic, stat = self._zookeeper.get("%s/%s" % (self._consumer_id_path, id_))
if topic == self._topic.name:
participants.append(id_)
except NoNodeException:
pass # disappeared between ``get_children`` and ``get``
participants = sorted(participants)
return participants
def _build_watch_callback(self, fn, proxy):
"""Return a function that's safe to use as a ChildrenWatch callback
Fixes the issue from https://github.com/Parsely/pykafka/issues/345
"""
def _callback(children):
# discover whether the referenced object still exists
try:
proxy.__repr__()
except ReferenceError:
return False
return fn(proxy, children)
return _callback
def _set_watches(self):
"""Set watches in zookeeper that will trigger rebalances.
Rebalances should be triggered whenever a broker, topic, or consumer
znode is changed in zookeeper. This ensures that the balance of the
consumer group remains up-to-date with the current state of the
cluster.
"""
proxy = weakref.proxy(self)
_brokers_changed = self._build_watch_callback(BalancedConsumer._brokers_changed, proxy)
_topics_changed = self._build_watch_callback(BalancedConsumer._topics_changed, proxy)
_consumers_changed = self._build_watch_callback(BalancedConsumer._consumers_changed, proxy)
self._setting_watches = True
# Set all our watches and then rebalance
broker_path = '/brokers/ids'
try:
self._broker_watcher = ChildrenWatch(
self._zookeeper, broker_path,
_brokers_changed
)
except NoNodeException:
raise Exception(
'The broker_path "%s" does not exist in your '
'ZooKeeper cluster -- is your Kafka cluster running?'
% broker_path)
self._topics_watcher = ChildrenWatch(
self._zookeeper,
'/brokers/topics',
_topics_changed
)
self._consumer_watcher = ChildrenWatch(
self._zookeeper, self._consumer_id_path,
_consumers_changed
)
self._setting_watches = False
def _add_self(self):
"""Register this consumer in zookeeper.
This method ensures that the number of participants is at most the
number of partitions.
"""
participants = self._get_participants()
if len(self._topic.partitions) <= len(participants):
raise KafkaException("Cannot add consumer: more consumers than partitions")
self._zookeeper.create(
self._path_self, self._topic.name, ephemeral=True, makepath=True)
@property
def _path_self(self):
"""Path where this consumer should be registered in zookeeper"""
return '{path}/{id_}'.format(
path=self._consumer_id_path,
id_=self._consumer_id
)
def _rebalance(self):
"""Claim partitions for this consumer.
This method is called whenever a zookeeper watch is triggered.
"""
if self._consumer is not None:
self.commit_offsets()
# this is necessary because we can't stop() while the lock is held
# (it's not an RLock)
should_stop = False
with self._rebalancing_lock:
if not self._running:
raise ConsumerStoppedException
log.info('Rebalancing consumer %s for topic %s.' % (
self._consumer_id, self._topic.name)
)
for i in range(self._rebalance_max_retries):
try:
# If retrying, be sure to make sure the
# partition allocation is correct.
participants = self._get_participants()
if self._consumer_id not in participants:
# situation that only occurs if our zk session expired
self._add_self()
participants.append(self._consumer_id)
new_partitions = self._decide_partitions(participants)
if not new_partitions:
should_stop = True
log.warning("No partitions assigned to consumer %s - stopping",
self._consumer_id)
break
# Update zk with any changes:
# Note that we explicitly fetch our set of held partitions
# from zk, rather than assuming it will be identical to
# `self.partitions`. This covers the (rare) situation
# where due to an interrupted connection our zk session
# has expired, in which case we'd hold zero partitions on
# zk, but `self._partitions` may be outdated and non-empty
current_zk_parts = self._get_held_partitions()
self._remove_partitions(current_zk_parts - new_partitions)
self._add_partitions(new_partitions - current_zk_parts)
# Only re-create internal consumer if something changed.
if new_partitions != self._partitions:
cns = self._get_internal_consumer(list(new_partitions))
if self._post_rebalance_callback is not None:
old_offsets = (self._consumer.held_offsets
if self._consumer else dict())
new_offsets = cns.held_offsets
reset_offsets = self._post_rebalance_callback(
self, old_offsets, new_offsets)
if reset_offsets:
cns.reset_offsets(partition_offsets=[
(cns.partitions[id_], offset) for
(id_, offset) in iteritems(reset_offsets)])
self._consumer = cns
log.info('Rebalancing Complete.')
break
except PartitionOwnedError as ex:
if i == self._rebalance_max_retries - 1:
log.warning('Failed to acquire partition %s after %d retries.',
ex.partition, i)
raise
log.info('Unable to acquire partition %s. Retrying', ex.partition)
time.sleep(i * (self._rebalance_backoff_ms / 1000))
if should_stop:
self.stop()
def _path_from_partition(self, p):
"""Given a partition, return its path in zookeeper.
:type p: :class:`pykafka.partition.Partition`
"""
return "%s/%s-%s" % (self._topic_path, p.leader.id, p.id)
def _remove_partitions(self, partitions):
"""Remove partitions from the zookeeper registry for this consumer.
:param partitions: The partitions to remove.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
# TODO pass zk node version to make sure we still own this node
self._zookeeper.delete(self._path_from_partition(p))
def _add_partitions(self, partitions):
"""Add partitions to the zookeeper registry for this consumer.
:param partitions: The partitions to add.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
try:
self._zookeeper.create(
self._path_from_partition(p),
value=get_bytes(self._consumer_id),
ephemeral=True
)
except NodeExistsError:
raise PartitionOwnedError(p)
def _get_held_partitions(self):
"""Build a set of partitions zookeeper says we own"""
zk_partition_ids = set()
all_partitions = self._zookeeper.get_children(self._topic_path)
for partition_slug in all_partitions:
try:
owner_id, stat = self._zookeeper.get(
'{path}/{slug}'.format(
path=self._topic_path, slug=partition_slug))
if owner_id == get_bytes(self._consumer_id):
zk_partition_ids.add(int(partition_slug.split('-')[1]))
except NoNodeException:
pass # disappeared between ``get_children`` and ``get``
return set(self._topic.partitions[_id] for _id in zk_partition_ids)
def _check_held_partitions(self):
"""Double-check held partitions against zookeeper
True if the partitions held by this consumer are the ones that
zookeeper thinks it's holding, else False.
"""
log.info("Checking held partitions against ZooKeeper")
zk_partitions = self._get_held_partitions()
if zk_partitions != self._partitions:
log.warning("Internal partition registry doesn't match ZooKeeper!")
log.debug("Internal partition ids: %s\nZooKeeper partition ids: %s",
self._partitions, zk_partitions)
return False
return True
@_catch_thread_exception
def _brokers_changed(self, brokers):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by broker change ({})".format(
self._consumer_id))
self._rebalance()
@_catch_thread_exception
def _consumers_changed(self, consumers):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by consumer change ({})".format(
self._consumer_id))
self._rebalance()
@_catch_thread_exception
def _topics_changed(self, topics):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by topic change ({})".format(
self._consumer_id))
self._rebalance()
def reset_offsets(self, partition_offsets=None):
"""Reset offsets for the specified partitions
Issue an OffsetRequest for each partition and set the appropriate
returned offset in the OwnedPartition
:param partition_offsets: (`partition`, `offset`) pairs to reset
where `partition` is the partition for which to reset the offset
and `offset` is the new offset the partition should have
:type partition_offsets: Iterable of
(:class:`pykafka.partition.Partition`, int)
"""
self._raise_worker_exceptions()
if not self._consumer:
raise ConsumerStoppedException("Internal consumer is stopped")
self._consumer.reset_offsets(partition_offsets=partition_offsets)
def consume(self, block=True):
"""Get one message from the consumer
:param block: Whether to block while waiting for a message
:type block: bool
"""
def consumer_timed_out():
"""Indicates whether the consumer has received messages recently"""
if self._consumer_timeout_ms == -1:
return False
disp = (time.time() - self._last_message_time) * 1000.0
return disp > self._consumer_timeout_ms
if not self._partitions:
raise NoPartitionsForConsumerException()
message = None
self._last_message_time = time.time()
while message is None and not consumer_timed_out():
self._raise_worker_exceptions()
try:
message = self._consumer.consume(block=block)
except ConsumerStoppedException:
if not self._running:
raise
continue
if message:
self._last_message_time = time.time()
if not block:
return message
return message
def __iter__(self):
"""Yield an infinite stream of messages until the consumer times out"""
while True:
message = self.consume(block=True)
if not message:
raise StopIteration
yield message
def commit_offsets(self):
"""Commit offsets for this consumer's partitions
Uses the offset commit/fetch API
"""
self._raise_worker_exceptions()
return self._consumer.commit_offsets()
class SimpleSwitch13(app_manager.RyuApp):
OFP_VERSIONS = [ofproto_v1_3.OFP_VERSION]
def __init__(self, *args, **kwargs):
super(SimpleSwitch13, self).__init__(*args, **kwargs)
self.mac_to_port = {}
self.zkConf = {'root':'/multicontroller', 'topo':'/topology',
'swstat':'/swstat', 'counter': '/counter'}
self.zk = KazooClient('127.0.0.1:2181')
self.zk.start()
self.ip = '202.201.3.51'
self.sws = {}
self.gid = random.randint(0, 10000)
self.dps = {}
self.links = []
self.interval = 5
self.role = OFPCR_ROLE_EQUAL
self.topoThread = hub.spawn(self._topoThread)
self.linkThread = hub.spawn(self._linkDiscover)
self.clearLinkThread = hub.spawn(self._cleanLinks)
self.clearLinkThread = hub.spawn(self._cleanSwitches)
def _cleanSwitches(self):
while True:
self.sws = {k:self.sws[k] for k in self.sws if self.sws[k]}
hub.sleep(self.interval)
def _topoThread(self):
linkNode = self.zkConf['root'] + self.zkConf['topo'] + self.ip
if self.zk.exists(linkNode):
self.zk.set(linkNode, json.dumps(self.links))
else:
self.zk.create(linkNode, json.dumps(self.links))
hub.sleep(self.interval)
def _linkDiscover(self):
while True:
for dpid in self.dps:
self.sendSlldp(dpid)
hub.sleep(self.interval)
def sendSlldp(self, dpid):
dp = self.dps.get(dpid)
if dp is None:
return
actions = [dp.ofproto_parser.OFPActionOutput(dp.ofproto.OFPP_FLOOD)]
pkt = packet.Packet()
pkt.add_protocol(ethernet.ethernet(ethertype=ETH_TYPE_SLLDP,
dst=SLLDP_MAC_DST, src=SLLDP_MAC_SRC))
pkt.add_protocol(slldp(dp.id))
pkt.serialize()
slldpPacket = pkt.data
out = dp.ofproto_parser.OFPPacketOut(
datapath=dp, in_port=dp.ofproto.OFPP_CONTROLLER,
buffer_id=dp.ofproto.OFP_NO_BUFFER, actions=actions,
data=slldpPacket)
dp.send_msg(out)
def getLinks(self):
topoNode = self.zkConf['root'] + self.zkConf['topo']
ips = self.zk.get_children(topoNode)
res = []
for ip in ips:
links = self.zk.get(topoNode + ip)[0]
for link in links:
res.append(link)
return res
@set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
def _packet_in_handler(self, ev):
msg = ev.msg
pkt = packet.Packet(msg.data)
eth = pkt.get_protocols(ethernet.ethernet)[0]
dst = eth.dst
# SLLDP packet
if dst == SLLDP_MAC_DST:
self.handleSlldp(ev)
return
# process packet_in message in subclass
self.packet_in_process(ev)
def handleSlldp(self, ev):
msg = ev.msg
datapath = msg.datapath
dpid = datapath.id
inPort = msg.match['in_port']
pkt = packet.Packet(msg.data)
slldpBuff = pkt.get_protocols(ethernet.ethernet)[2]
dpidSrc, _ = slldp.parser(slldpBuff)
self.links.append({'srcdpid': dpidSrc,
'dst': {'dpid': dpid, 'port': inPort},
'time': time.time()})
def _cleanLinks(self):
while True:
now = time.time()
self.links = [l for l in self.links if now - l['time'] < self.interval]
hub.sleep(self.interval)
@abc.abstractmethod
def packet_in_process(self, ev):
pass
@set_ev_cls(event.EventSwitchEnter)
def switch_enter(self, ev):
dpid = ev.datapath.id
self.sws[dpid] = True
self.dps[dpid] = ev.datapath
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
self.zk.ensure_path(dpNode)
if self.election(dpid):
self.role = OFPCR_ROLE_MASTER
else:
self.role = OFPCR_ROLE_SLAVE
self.countUp(dpid)
self.roleRequest(dpid, self.role)
mflag = dpNode + '/' + 'master'
DataWatch(self.zk, mflag, self.masterWatcher)
def masterWatcher(self, data, stat, ev):
if ev and ev.type == 'DELETED':
_, _, dpid, _ = ev.path.split('/')
dpid = str_to_dpid(dpid)
if self.sws.get(dpid):
if self.election(dpid):
self.role = OFPCR_ROLE_MASTER
self.roleRequest(dpid, self.role)
return self.sws.get(dpid)
def election(self, dpid):
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
mflag = dpNode + '/' + 'master'
while not self.zk.exists(mflag):
mlock = self.zk.Lock(dpNode + '/' + 'mlock', self.ip)
with mlock:
if not self.zk.exists(mflag):
self.zk.create(mflag, self.ip, ephemeral=True)
if self.zk.exists(mflag):
if self.zk.get(mflag) == self.ip:
return True
else:
return False
else:
time.sleep(random.randint(0, 100)/500.0)
return False
def roleRequest(self, dp, role):
msg = dp.ofproto_parser.OFPRoleRequest(dp, role, self.gid)
dp.send_msg(msg)
def getCount(self, dpid):
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
countNode = dpNode + self.zkConf['counter']
counters = self.zk.get_children(countNode)
return len(counters)
def countUp(self, dpid):
countNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid) + self.zkConf['counter']
self.zk.ensure_path(countNode)
self.zk.create(countNode+uuid4().hex, 'alive', ephemeral=True)
@set_ev_cls(event.EventSwitchLeave)
def switch_levave(self, ev):
dpid = ev.datapath
count = self.getCount(dpid)
self.sws[dpid] = False
if count == 0:
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
self.zk.delete(dpNode, recursive=True)
children = zk.get_children('/app')
printout("[INIT]", WHITE)
print(children)
portno = 0
portnum = 0
ranges = []
serverDied = False
if len(children) == 1:
portno = 8080
printout("[MASTER]", RED)
print("MASTER INIT")
if zk.exists('/meta'):
zk.delete('/meta', recursive=True)
zk.create('/meta/master', b'8080', makepath=True)
zk.create('/meta/status', 'INIT'.encode("utf-8"), makepath=True)
zk.create('/meta/lastport', '8080'.encode('utf-8'), makepath=True)
time.sleep(5)
children = zk.get_children('/app')
config = {
"mapper": {},
"lastDead": {
"backup": -1,
# backup is the server port whose data the new server needs to retrieve
"portno": -1
# dead server portno
},
"numOfServers": len(children)
while True:
print("cnt(%s)************************************************" % (cnt))
for child in redis_list:
child_path = root + '/' + child
if child in conn:
rs = conn[host]
else:
host, port = child.split(':')
rs = redis.Redis(host=host, port=port)
conn[host] = rs
transaction = zk.transaction()
try:
rs.ping()
except ConnectionError:
if zk.exists(child_path) != None:
zk.delete(child_path)
print("%s: error" % (child_path))
else:
if zk.exists(child_path) == None:
zk.create(child_path)
print("%s: ok" % (child_path))
transaction.commit()
time.sleep(10)
cnt += 1
zk.stop()
from kazoo.client import KazooClient
import sys
if __name__ == "__main__":
kz = KazooClient('127.0.0.1:2181')
kz.start()
if len(sys.argv) >= 2 and sys.argv[1] == "clear":
kz.ensure_path('/dimint/overlord/host_list')
kz.delete('/dimint/overlord/host_list', recursive=True)
kz.ensure_path('/dimint/overlord/host_list')
kz.ensure_path('/dimint/node/list')
kz.delete('/dimint/node/list', recursive=True)
kz.ensure_path('/dimint/node/list')
kz.ensure_path('/dimint/node/role')
kz.delete('/dimint/node/role', recursive=True)
kz.ensure_path('/dimint/node/role')
else:
kz.ensure_path('/dimint/overlord/host_list')
print('overlord list : {0}'.format(
kz.get_children('/dimint/overlord/host_list')))
kz.ensure_path('/dimint/node/list')
print('node list : {0}'.format(kz.get_children('/dimint/node/list')))
kz.ensure_path('/dimint/node/role')
for master in kz.get_children('/dimint/node/role'):
print('master node : {0}'.format(master))
for slave in kz.get_children(
'/dimint/node/role/{0}'.format(master)):
print("\tslave node : {0}".format(slave))
kz.stop()
client = docker.from_env()
# A set of global variables to handle entering various sections of the code
not_called_by_scale = True
deleted_master = False
# Setting the initial active containers pid and uuid and the counting the number or requests
active_containers = {}
counter = 0
# Establish a connection with the Zookeeper
zk = KazooClient(hosts='10.0.2.3:2181')
zk.start()
# Create a parent node for the child container if one doesn't exist
if (zk.exists("/producer")):
zk.delete("/producer", recursive=True)
zk.ensure_path("/producer")
# Run the initial two workers of master and slave on the cry_cloud network
# The time.sleep argument is for the docker service to have enough time to create a new container
temp1 = client.containers.run("worker",
detach=True,
auto_remove=True,
network="cry_cloud")
time.sleep(10)
temp = client.containers.run("worker",
detach=True,
auto_remove=True,
network="cry_cloud")
time.sleep(10)
class ZkStateManager(StateManager):
"""
State manager which connects to zookeeper and
gets and sets states from there.
"""
def __init__(self, name, host, port, rootpath, tunnelhost):
self.name = name
self.host = host
self.port = port
self.tunnelhost = tunnelhost
self.rootpath = rootpath
def start(self):
""" state Zookeeper """
if self.is_host_port_reachable():
self.client = KazooClient(self.hostport)
else:
localport = self.establish_ssh_tunnel()
self.client = KazooClient("localhost:" + str(localport))
self.client.start()
def on_connection_change(state):
""" callback to log """
LOG.info("Connection state changed to: " + state)
self.client.add_listener(on_connection_change)
def stop(self):
""" stop Zookeeper """
self.client.stop()
self.terminate_ssh_tunnel()
# pylint: disable=function-redefined
def get_topologies(self, callback=None):
""" get topologies """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""Custom callback to get the topologies right now."""
ret["result"] = data
try:
self._get_topologies_with_watch(callback, isWatching)
except NoNodeError:
self.client.stop()
path = self.get_topologies_path()
raise StateException("Error required topology path '%s' not found" % (path),
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
# The topologies are now populated with the data.
return ret["result"]
def _get_topologies_with_watch(self, callback, isWatching):
"""
Helper function to get topologies with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_topologies_path()
if isWatching:
LOG.info("Adding children watch for path: " + path)
# pylint: disable=unused-variable
@self.client.ChildrenWatch(path)
def watch_topologies(topologies):
""" callback to watch topologies """
callback(topologies)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def get_topology(self, topologyName, callback=None):
""" get topologies """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""Custom callback to get the topologies right now."""
ret["result"] = data
self._get_topology_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_topology_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_topology_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_topology(data, stats):
""" watch topology """
if data:
topology = Topology()
topology.ParseFromString(data)
callback(topology)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_topology(self, topologyName, topology):
""" crate topology """
if not topology or not topology.IsInitialized():
raise StateException("Topology protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_topology_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
topologyString = topology.SerializeToString()
try:
self.client.create(path, value=topologyString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating topology",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating topology",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating topology",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_topology(self, topologyName):
""" delete topology """
path = self.get_topology_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deteling topology",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting topology",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting topology",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_packing_plan(self, topologyName, callback=None):
""" get packing plan """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
""" Custom callback to get the topologies right now. """
ret["result"] = data
self._get_packing_plan_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_packing_plan_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get packing_plan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_packing_plan_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-argument,unused-variable
@self.client.DataWatch(path)
def watch_packing_plan(data, stats):
""" watch the packing plan for updates """
if data:
packing_plan = PackingPlan()
packing_plan.ParseFromString(data)
callback(packing_plan)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def get_pplan(self, topologyName, callback=None):
""" get physical plan """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_pplan_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_pplan_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_pplan_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_pplan(data, stats):
""" invoke callback to watch physical plan """
if data:
pplan = PhysicalPlan()
pplan.ParseFromString(data)
callback(pplan)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_pplan(self, topologyName, pplan):
""" create physical plan """
if not pplan or not pplan.IsInitialized():
raise StateException("Physical Plan protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_pplan_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
pplanString = pplan.SerializeToString()
try:
self.client.create(path, value=pplanString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating pplan",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating pplan",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating pplan",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_pplan(self, topologyName):
""" delete physical plan info """
path = self.get_pplan_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deleting pplan",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting pplan",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting pplan",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_execution_state(self, topologyName, callback=None):
""" get execution state """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_execution_state_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_execution_state_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get execution state with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_execution_state_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_execution_state(data, stats):
""" invoke callback to watch execute state """
if data:
executionState = ExecutionState()
executionState.ParseFromString(data)
callback(executionState)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_execution_state(self, topologyName, executionState):
""" create execution state """
if not executionState or not executionState.IsInitialized():
raise StateException("Execution State protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_execution_state_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
executionStateString = executionState.SerializeToString()
try:
self.client.create(path, value=executionStateString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating execution state",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating execution state",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating execution state",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_execution_state(self, topologyName):
""" delete execution state """
path = self.get_execution_state_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deleting execution state",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting execution state",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting execution state",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_tmaster(self, topologyName, callback=None):
""" get tmaster """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_tmaster_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_tmaster_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_tmaster_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_tmaster(data, stats):
""" invoke callback to watch tmaster """
if data:
tmaster = TMasterLocation()
tmaster.ParseFromString(data)
callback(tmaster)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def get_scheduler_location(self, topologyName, callback=None):
""" get scheduler location """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the scheduler location right now.
"""
ret["result"] = data
self._get_scheduler_location_with_watch(topologyName, callback, isWatching)
return ret["result"]
def _get_scheduler_location_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get scheduler location with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_scheduler_location_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_scheduler_location(data, stats):
""" invoke callback to watch scheduler location """
if data:
scheduler_location = SchedulerLocation()
scheduler_location.ParseFromString(data)
callback(scheduler_location)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def start_kafka(zk_client_port, broker_listen_port, broker_id=0):
if not os.path.exists(kafka_bdir):
output, _ = call_command_("mkdir " + kafka_bdir)
kafka_download = 'wget -nv --tries=3 -c -O ' + kafka_bdir + kafka_dl + \
' https://github.com/Juniper/contrail-third-party-cache/blob/master/kafka' + \
kafka_dl + '?raw=true'
if not os.path.exists(kafka_bdir + kafka_dl):
process = subprocess.Popen(kafka_download.split(' '))
process.wait()
if process.returncode is not 0:
return False
basefile = kafka_version
kafkabase = "/tmp/kafka.%s.%d/" % (os.getenv('USER',
'None'), broker_listen_port)
confdir = kafkabase + basefile + "/config/"
output, _ = call_command_("rm -rf " + kafkabase)
output, _ = call_command_("mkdir " + kafkabase)
logging.info('Check zookeeper in %d' % zk_client_port)
zk = KazooClient(hosts='127.0.0.1:' + str(zk_client_port), timeout=60.0)
try:
zk.start()
zk.delete("/brokers", recursive=True)
zk.delete("/consumers", recursive=True)
zk.delete("/controller", recursive=True)
except:
logging.info("Zookeeper client cannot connect")
zk.stop()
return False
zk.stop()
logging.info('Installing kafka in ' + kafkabase)
x = os.system("cat " + kafka_bdir + kafka_dl + " | tar -xpzf - -C " +
kafkabase)
if 0 != x:
logging.error("Cannot install kafka")
return False
logging.info('kafka Port %d' % broker_listen_port)
replace_string_(confdir + "server.properties",
[("#listeners=PLAINTEXT://:9092",
"listeners=PLAINTEXT://:" + str(broker_listen_port))])
#Replace the brokerid and port # in the config file
replace_string_(
confdir + "server.properties",
[("broker.id=0", "broker.id=" + str(broker_id)),
("zookeeper.connect=localhost:2181",
"zookeeper.connect=localhost:%d" % zk_client_port),
("log.dirs=/tmp/kafka-logs", "log.dirs=" + kafkabase + "logs")])
replace_string_(kafkabase + basefile + "/bin/kafka-server-stop.sh",
[("grep -v grep", "grep %s | grep -v grep" % kafkabase)])
replace_string_(kafkabase + basefile + "/bin/kafka-server-stop.sh",
[("SIGINT", "SIGKILL")])
replace_string_(kafkabase + basefile + "/bin/kafka-server-stop.sh",
[("#!/bin/sh", "#!/bin/sh -x")])
output, _ = call_command_("chmod +x " + kafkabase + basefile +
"/bin/kafka-server-stop.sh")
# Extra options for JMX : -Djava.net.preferIPv4Stack=true -Djava.rmi.server.hostname=xx.xx.xx.xx
output, _ = call_command_(kafkabase + basefile +
"/bin/kafka-server-start.sh -daemon " +
kafkabase + basefile +
"/config/server.properties")
count = 0
start_wait = os.getenv('CONTRIAL_ANALYTICS_TEST_MAX_START_WAIT_TIME', 15)
while count < start_wait:
try:
logging.info('Trying to connect...')
kk = KafkaClient("localhost:%d" % broker_listen_port)
except:
count += 1
time.sleep(1)
else:
return True
logging.info("Kafka client cannot connect. Kafka logfile below:")
with open(kafkabase + basefile + "/logs/kafkaServer.out", 'r') as fin:
logging.info(fin.read())
return False
class KazooCommandProxy():
def __init__(self, module):
self.module = module
self.zk = KazooClient(module.params['hosts'])
def absent(self):
return self._absent(self.module.params['name'])
def exists(self, znode):
return self.zk.exists(znode)
def list(self):
children = self.zk.get_children(self.module.params['name'])
return True, {'count': len(children), 'items': children, 'msg': 'Retrieved znodes in path.',
'znode': self.module.params['name']}
def present(self):
return self._present(self.module.params['name'], self.module.params['value'])
def get(self):
return self._get(self.module.params['name'])
def shutdown(self):
self.zk.stop()
self.zk.close()
def start(self):
self.zk.start()
def wait(self):
return self._wait(self.module.params['name'], self.module.params['timeout'])
def _absent(self, znode):
if self.exists(znode):
self.zk.delete(znode, recursive=self.module.params['recursive'])
return True, {'changed': True, 'msg': 'The znode was deleted.'}
else:
return True, {'changed': False, 'msg': 'The znode does not exist.'}
def _get(self, path):
if self.exists(path):
value, zstat = self.zk.get(path)
stat_dict = {}
for i in dir(zstat):
if not i.startswith('_'):
attr = getattr(zstat, i)
if isinstance(attr, (int, str)):
stat_dict[i] = attr
result = True, {'msg': 'The node was retrieved.', 'znode': path, 'value': value,
'stat': stat_dict}
else:
result = False, {'msg': 'The requested node does not exist.'}
return result
def _present(self, path, value):
if self.exists(path):
(current_value, zstat) = self.zk.get(path)
if value != current_value:
self.zk.set(path, value)
return True, {'changed': True, 'msg': 'Updated the znode value.', 'znode': path,
'value': value}
else:
return True, {'changed': False, 'msg': 'No changes were necessary.', 'znode': path, 'value': value}
else:
self.zk.create(path, value, makepath=True)
return True, {'changed': True, 'msg': 'Created a new znode.', 'znode': path, 'value': value}
def _wait(self, path, timeout, interval=5):
lim = time.time() + timeout
while time.time() < lim:
if self.exists(path):
return True, {'msg': 'The node appeared before the configured timeout.',
'znode': path, 'timeout': timeout}
else:
time.sleep(interval)
return False, {'msg': 'The node did not appear before the operation timed out.', 'timeout': timeout,
'znode': path}
class BalancedConsumer():
"""
A self-balancing consumer for Kafka that uses ZooKeeper to communicate
with other balancing consumers.
Maintains a single instance of SimpleConsumer, periodically using the
consumer rebalancing algorithm to reassign partitions to this
SimpleConsumer.
"""
def __init__(self,
topic,
cluster,
consumer_group,
fetch_message_max_bytes=1024 * 1024,
num_consumer_fetchers=1,
auto_commit_enable=False,
auto_commit_interval_ms=60 * 1000,
queued_max_messages=2000,
fetch_min_bytes=1,
fetch_wait_max_ms=100,
offsets_channel_backoff_ms=1000,
offsets_commit_max_retries=5,
auto_offset_reset=OffsetType.EARLIEST,
consumer_timeout_ms=-1,
rebalance_max_retries=5,
rebalance_backoff_ms=2 * 1000,
zookeeper_connection_timeout_ms=6 * 1000,
zookeeper_connect='127.0.0.1:2181',
zookeeper=None,
auto_start=True,
reset_offset_on_start=False):
"""Create a BalancedConsumer instance
:param topic: The topic this consumer should consume
:type topic: :class:`pykafka.topic.Topic`
:param cluster: The cluster to which this consumer should connect
:type cluster: :class:`pykafka.cluster.Cluster`
:param consumer_group: The name of the consumer group this consumer
should join.
:type consumer_group: bytes
:param fetch_message_max_bytes: The number of bytes of messages to
attempt to fetch with each fetch request
:type fetch_message_max_bytes: int
:param num_consumer_fetchers: The number of workers used to make
FetchRequests
:type num_consumer_fetchers: int
:param auto_commit_enable: If true, periodically commit to kafka the
offset of messages already fetched by this consumer. This also
requires that `consumer_group` is not `None`.
:type auto_commit_enable: bool
:param auto_commit_interval_ms: The frequency (in milliseconds) at which
the consumer's offsets are committed to kafka. This setting is
ignored if `auto_commit_enable` is `False`.
:type auto_commit_interval_ms: int
:param queued_max_messages: The maximum number of messages buffered for
consumption in the internal
:class:`pykafka.simpleconsumer.SimpleConsumer`
:type queued_max_messages: int
:param fetch_min_bytes: The minimum amount of data (in bytes) that the
server should return for a fetch request. If insufficient data is
available, the request will block until sufficient data is available.
:type fetch_min_bytes: int
:param fetch_wait_max_ms: The maximum amount of time (in milliseconds)
that the server will block before answering a fetch request if
there isn't sufficient data to immediately satisfy `fetch_min_bytes`.
:type fetch_wait_max_ms: int
:param offsets_channel_backoff_ms: Backoff time to retry failed offset
commits and fetches.
:type offsets_channel_backoff_ms: int
:param offsets_commit_max_retries: The number of times the offset commit
worker should retry before raising an error.
:type offsets_commit_max_retries: int
:param auto_offset_reset: What to do if an offset is out of range. This
setting indicates how to reset the consumer's internal offset
counter when an `OffsetOutOfRangeError` is encountered.
:type auto_offset_reset: :class:`pykafka.common.OffsetType`
:param consumer_timeout_ms: Amount of time (in milliseconds) the
consumer may spend without messages available for consumption
before returning None.
:type consumer_timeout_ms: int
:param rebalance_max_retries: The number of times the rebalance should
retry before raising an error.
:type rebalance_max_retries: int
:param rebalance_backoff_ms: Backoff time (in milliseconds) between
retries during rebalance.
:type rebalance_backoff_ms: int
:param zookeeper_connection_timeout_ms: The maximum time (in
milliseconds) that the consumer waits while establishing a
connection to zookeeper.
:type zookeeper_connection_timeout_ms: int
:param zookeeper_connect: Comma-separated (ip1:port1,ip2:port2) strings
indicating the zookeeper nodes to which to connect.
:type zookeeper_connect: str
:param zookeeper: A KazooClient connected to a Zookeeper instance.
If provided, `zookeeper_connect` is ignored.
:type zookeeper: :class:`kazoo.client.KazooClient`
:param auto_start: Whether the consumer should begin communicating
with zookeeper after __init__ is complete. If false, communication
can be started with `start()`.
:type auto_start: bool
:param reset_offset_on_start: Whether the consumer should reset its
internal offset counter to `self._auto_offset_reset` and commit that
offset immediately upon starting up
:type reset_offset_on_start: bool
"""
self._cluster = cluster
self._consumer_group = consumer_group
self._topic = topic
self._auto_commit_enable = auto_commit_enable
self._auto_commit_interval_ms = auto_commit_interval_ms
self._fetch_message_max_bytes = fetch_message_max_bytes
self._fetch_min_bytes = fetch_min_bytes
self._rebalance_max_retries = rebalance_max_retries
self._num_consumer_fetchers = num_consumer_fetchers
self._queued_max_messages = queued_max_messages
self._fetch_wait_max_ms = fetch_wait_max_ms
self._rebalance_backoff_ms = rebalance_backoff_ms
self._consumer_timeout_ms = consumer_timeout_ms
self._offsets_channel_backoff_ms = offsets_channel_backoff_ms
self._offsets_commit_max_retries = offsets_commit_max_retries
self._auto_offset_reset = auto_offset_reset
self._zookeeper_connect = zookeeper_connect
self._zookeeper_connection_timeout_ms = zookeeper_connection_timeout_ms
self._reset_offset_on_start = reset_offset_on_start
self._running = False
self._rebalancing_lock = cluster.handler.Lock()
self._consumer = None
self._consumer_id = "{hostname}:{uuid}".format(
hostname=socket.gethostname(),
uuid=uuid4()
)
self._partitions = set()
self._setting_watches = True
self._topic_path = '/consumers/{group}/owners/{topic}'.format(
group=self._consumer_group,
topic=self._topic.name)
self._consumer_id_path = '/consumers/{group}/ids'.format(
group=self._consumer_group)
self._zookeeper = None
if zookeeper is not None:
self._zookeeper = zookeeper
if auto_start is True:
self.start()
def __repr__(self):
return "<{module}.{name} at {id_} (consumer_group={group})>".format(
module=self.__class__.__module__,
name=self.__class__.__name__,
id_=hex(id(self)),
group=self._consumer_group
)
def _setup_checker_worker(self):
"""Start the zookeeper partition checker thread"""
def checker():
while True:
time.sleep(120)
if not self._running:
break
self._check_held_partitions()
log.debug("Checker thread exiting")
log.debug("Starting checker thread")
return self._cluster.handler.spawn(checker)
@property
def partitions(self):
return self._consumer.partitions if self._consumer else None
@property
def held_offsets(self):
"""Return a map from partition id to held offset for each partition"""
if not self._consumer:
return None
return dict((p.partition.id, p.last_offset_consumed)
for p in self._consumer._partitions_by_id.itervalues())
def start(self):
"""Open connections and join a cluster."""
if self._zookeeper is None:
self._setup_zookeeper(self._zookeeper_connect,
self._zookeeper_connection_timeout_ms)
self._zookeeper.ensure_path(self._topic_path)
self._add_self()
self._set_watches()
self._rebalance()
self._running = True
self._setup_checker_worker()
def stop(self):
"""Close the zookeeper connection and stop consuming.
This method should be called as part of a graceful shutdown process.
"""
self._zookeeper.stop()
self._consumer.stop()
self._running = False
def _setup_zookeeper(self, zookeeper_connect, timeout):
"""Open a connection to a ZooKeeper host.
:param zookeeper_connect: The 'ip:port' address of the zookeeper node to
which to connect.
:type zookeeper_connect: str
:param timeout: Connection timeout (in milliseconds)
:type timeout: int
"""
self._zookeeper = KazooClient(zookeeper_connect, timeout=timeout / 1000)
self._zookeeper.start()
def _setup_internal_consumer(self, start=True):
"""Instantiate an internal SimpleConsumer.
If there is already a SimpleConsumer instance held by this object,
disable its workers and mark it for garbage collection before
creating a new one.
"""
reset_offset_on_start = self._reset_offset_on_start
if self._consumer is not None:
self._consumer.stop()
# only use this setting for the first call to
# _setup_internal_consumer. subsequent calls should not
# reset the offsets, since they can happen at any time
reset_offset_on_start = False
self._consumer = SimpleConsumer(
self._topic,
self._cluster,
consumer_group=self._consumer_group,
partitions=list(self._partitions),
auto_commit_enable=self._auto_commit_enable,
auto_commit_interval_ms=self._auto_commit_interval_ms,
fetch_message_max_bytes=self._fetch_message_max_bytes,
fetch_min_bytes=self._fetch_min_bytes,
num_consumer_fetchers=self._num_consumer_fetchers,
queued_max_messages=self._queued_max_messages,
fetch_wait_max_ms=self._fetch_wait_max_ms,
consumer_timeout_ms=self._consumer_timeout_ms,
offsets_channel_backoff_ms=self._offsets_channel_backoff_ms,
offsets_commit_max_retries=self._offsets_commit_max_retries,
auto_offset_reset=self._auto_offset_reset,
reset_offset_on_start=reset_offset_on_start,
auto_start=start
)
def _decide_partitions(self, participants):
"""Decide which partitions belong to this consumer.
Uses the consumer rebalancing algorithm described here
http://kafka.apache.org/documentation.html
It is very important that the participants array is sorted,
since this algorithm runs on each consumer and indexes into the same
array. The same array index operation must return the same
result on each consumer.
:param participants: Sorted list of ids of all other consumers in this
consumer group.
:type participants: Iterable of `bytes`
"""
# Freeze and sort partitions so we always have the same results
p_to_str = lambda p: '-'.join([str(p.topic.name), str(p.leader.id), str(p.id)])
all_parts = self._topic.partitions.values()
all_parts = sorted(all_parts, key=p_to_str)
# get start point, # of partitions, and remainder
participants = sorted(participants) # just make sure it's sorted.
idx = participants.index(self._consumer_id)
parts_per_consumer = len(all_parts) // len(participants)
remainder_ppc = len(all_parts) % len(participants)
start = parts_per_consumer * idx + min(idx, remainder_ppc)
num_parts = parts_per_consumer + (0 if (idx + 1 > remainder_ppc) else 1)
# assign partitions from i*N to (i+1)*N - 1 to consumer Ci
new_partitions = itertools.islice(all_parts, int(start), int(start + num_parts))
new_partitions = set(new_partitions)
log.info('Balancing %i participants for %i partitions.\nOwning %i partitions.',
len(participants), len(all_parts), len(new_partitions))
log.debug('My partitions: %s', [p_to_str(p) for p in new_partitions])
return new_partitions
def _get_participants(self):
"""Use zookeeper to get the other consumers of this topic.
:return: A sorted list of the ids of the other consumers of this
consumer's topic
"""
try:
consumer_ids = self._zookeeper.get_children(self._consumer_id_path)
except NoNodeException:
log.debug("Consumer group doesn't exist. "
"No participants to find")
return []
participants = []
for id_ in consumer_ids:
try:
topic, stat = self._zookeeper.get("%s/%s" % (self._consumer_id_path, id_))
if topic == self._topic.name:
participants.append(id_)
except NoNodeException:
pass # disappeared between ``get_children`` and ``get``
participants = sorted(participants)
return participants
def _set_watches(self):
"""Set watches in zookeeper that will trigger rebalances.
Rebalances should be triggered whenever a broker, topic, or consumer
znode is changed in zookeeper. This ensures that the balance of the
consumer group remains up-to-date with the current state of the
cluster.
"""
self._setting_watches = True
# Set all our watches and then rebalance
broker_path = '/brokers/ids'
try:
self._broker_watcher = ChildrenWatch(
self._zookeeper, broker_path,
self._brokers_changed
)
except NoNodeException:
raise Exception(
'The broker_path "%s" does not exist in your '
'ZooKeeper cluster -- is your Kafka cluster running?'
% broker_path)
self._topics_watcher = ChildrenWatch(
self._zookeeper,
'/brokers/topics',
self._topics_changed
)
self._consumer_watcher = ChildrenWatch(
self._zookeeper, self._consumer_id_path,
self._consumers_changed
)
self._setting_watches = False
def _add_self(self):
"""Register this consumer in zookeeper.
This method ensures that the number of participants is at most the
number of partitions.
"""
participants = self._get_participants()
if len(self._topic.partitions) <= len(participants):
raise KafkaException("Cannot add consumer: more consumers than partitions")
path = '{path}/{id_}'.format(
path=self._consumer_id_path,
id_=self._consumer_id
)
self._zookeeper.create(
path, self._topic.name, ephemeral=True, makepath=True)
def _rebalance(self):
"""Claim partitions for this consumer.
This method is called whenever a zookeeper watch is triggered.
"""
if self._consumer is not None:
self.commit_offsets()
with self._rebalancing_lock:
log.info('Rebalancing consumer %s for topic %s.' % (
self._consumer_id, self._topic.name)
)
for i in range(self._rebalance_max_retries):
try:
# If retrying, be sure to make sure the
# partition allocation is correct.
participants = self._get_participants()
partitions = self._decide_partitions(participants)
old_partitions = self._partitions - partitions
self._remove_partitions(old_partitions)
new_partitions = partitions - self._partitions
self._add_partitions(new_partitions)
# Only re-create internal consumer if something changed.
if old_partitions or new_partitions:
self._setup_internal_consumer()
log.info('Rebalancing Complete.')
break
except PartitionOwnedError as ex:
if i == self._rebalance_max_retries - 1:
log.warning('Failed to acquire partition %s after %d retries.',
ex.partition, i)
raise
log.info('Unable to acquire partition %s. Retrying', ex.partition)
time.sleep(i * (self._rebalance_backoff_ms / 1000))
def _path_from_partition(self, p):
"""Given a partition, return its path in zookeeper.
:type p: :class:`pykafka.partition.Partition`
"""
return "%s/%s-%s" % (self._topic_path, p.leader.id, p.id)
def _remove_partitions(self, partitions):
"""Remove partitions from the zookeeper registry for this consumer.
Also remove these partitions from the consumer's internal
partition registry.
:param partitions: The partitions to remove.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
assert p in self._partitions
self._zookeeper.delete(self._path_from_partition(p))
self._partitions -= partitions
def _add_partitions(self, partitions):
"""Add partitions to the zookeeper registry for this consumer.
Also add these partitions to the consumer's internal partition registry.
:param partitions: The partitions to add.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
try:
self._zookeeper.create(
self._path_from_partition(p),
value=get_bytes(self._consumer_id),
ephemeral=True
)
self._partitions.add(p)
except NodeExistsError:
raise PartitionOwnedError(p)
def _check_held_partitions(self):
"""Double-check held partitions against zookeeper
Ensure that the partitions held by this consumer are the ones that
zookeeper thinks it's holding. If not, rebalance.
"""
log.info("Checking held partitions against ZooKeeper")
# build a set of partition ids zookeeper says we own
zk_partition_ids = set()
all_partitions = self._zookeeper.get_children(self._topic_path)
for partition_slug in all_partitions:
owner_id, stat = self._zookeeper.get(
'{path}/{slug}'.format(
path=self._topic_path, slug=partition_slug))
if owner_id == self._consumer_id:
zk_partition_ids.add(int(partition_slug.split('-')[1]))
# build a set of partition ids we think we own
internal_partition_ids = set([p.id for p in self._partitions])
# compare the two sets, rebalance if necessary
if internal_partition_ids != zk_partition_ids:
log.warning("Internal partition registry doesn't match ZooKeeper!")
log.debug("Internal partition ids: %s\nZooKeeper partition ids: %s",
internal_partition_ids, zk_partition_ids)
self._rebalance()
def _brokers_changed(self, brokers):
if self._setting_watches:
return
log.debug("Rebalance triggered by broker change")
self._rebalance()
def _consumers_changed(self, consumers):
if self._setting_watches:
return
log.debug("Rebalance triggered by consumer change")
self._rebalance()
def _topics_changed(self, topics):
if self._setting_watches:
return
log.debug("Rebalance triggered by topic change")
self._rebalance()
def reset_offsets(self, partition_offsets=None):
"""Reset offsets for the specified partitions
Issue an OffsetRequest for each partition and set the appropriate
returned offset in the OwnedPartition
:param partition_offsets: (`partition`, `offset`) pairs to reset
where `partition` is the partition for which to reset the offset
and `offset` is the new offset the partition should have
:type partition_offsets: Iterable of
(:class:`pykafka.partition.Partition`, int)
"""
if not self._consumer:
raise ConsumerStoppedException("Internal consumer is stopped")
self._consumer.reset_offsets(partition_offsets=partition_offsets)
def consume(self, block=True):
"""Get one message from the consumer
:param block: Whether to block while waiting for a message
:type block: bool
"""
def consumer_timed_out():
"""Indicates whether the consumer has received messages recently"""
if self._consumer_timeout_ms == -1:
return False
disp = (time.time() - self._last_message_time) * 1000.0
return disp > self._consumer_timeout_ms
message = None
self._last_message_time = time.time()
while message is None and not consumer_timed_out():
try:
message = self._consumer.consume(block=block)
except ConsumerStoppedException:
if not self._running:
return
continue
if message:
self._last_message_time = time.time()
if not block:
return message
return message
def commit_offsets(self):
"""Commit offsets for this consumer's partitions
Uses the offset commit/fetch API
"""
return self._consumer.commit_offsets()
class ZooKeeper(AbstractDCS):
def __init__(self, name, config):
super(ZooKeeper, self).__init__(name, config)
hosts = config.get('hosts', [])
if isinstance(hosts, list):
hosts = ','.join(hosts)
self.exhibitor = None
if 'exhibitor' in config:
exhibitor = config['exhibitor']
interval = exhibitor.get('poll_interval', 300)
self.exhibitor = ExhibitorEnsembleProvider(exhibitor['hosts'],
exhibitor['port'],
poll_interval=interval)
hosts = self.exhibitor.zookeeper_hosts
self.client = KazooClient(hosts=hosts,
timeout=(config.get('session_timeout', None)
or 30),
command_retry={
'deadline':
(config.get('reconnect_timeout', None)
or 10),
'max_delay':
1,
'max_tries':
-1
},
connection_retry={
'max_delay': 1,
'max_tries': -1
})
self.client.add_listener(self.session_listener)
self._my_member_data = None
self.fetch_cluster = True
self.last_leader_operation = 0
self.client.start(None)
def session_listener(self, state):
if state in [KazooState.SUSPENDED, KazooState.LOST]:
self.cluster_watcher(None)
def cluster_watcher(self, event):
self.fetch_cluster = True
self.event.set()
def get_node(self, key, watch=None):
try:
ret = self.client.get(key, watch)
return (ret[0].decode('utf-8'), ret[1])
except NoNodeError:
return None
@staticmethod
def member(name, value, znode):
return Member.from_node(znode.version, name, znode.ephemeralOwner,
value)
def get_children(self, key, watch=None):
try:
return self.client.get_children(key, watch)
except NoNodeError:
return []
def load_members(self):
members = []
for member in self.get_children(self.members_path,
self.cluster_watcher):
data = self.get_node(self.members_path + member)
if data is not None:
members.append(self.member(member, *data))
return members
def _inner_load_cluster(self):
self.fetch_cluster = False
self.event.clear()
nodes = set(
self.get_children(self.client_path(''), self.cluster_watcher))
if not nodes:
self.fetch_cluster = True
# get initialize flag
initialize = (self.get_node(self.initialize_path)
or [None])[0] if self._INITIALIZE in nodes else None
# get list of members
members = self.load_members() if self._MEMBERS[:-1] in nodes else []
# get leader
leader = self.get_node(
self.leader_path) if self._LEADER in nodes else None
if leader:
client_id = self.client.client_id
if leader[0] == self._name and client_id is not None and client_id[
0] != leader[1].ephemeralOwner:
logger.info(
'I am leader but not owner of the session. Removing leader node'
)
self.client.delete(self.leader_path)
leader = None
if leader:
member = Member(-1, leader[0], None, {})
member = ([m for m in members if m.name == leader[0]]
or [member])[0]
leader = Leader(leader[1].version, leader[1].ephemeralOwner,
member)
self.fetch_cluster = member.index == -1
# failover key
failover = self.get_node(
self.failover_path,
watch=self.cluster_watcher) if self._FAILOVER in nodes else None
if failover:
failover = Failover.from_node(failover[1].version, failover[0])
# get last leader operation
optime = self.get_node(
self.leader_optime_path
) if self._OPTIME in nodes and self.fetch_cluster else None
self.last_leader_operation = 0 if optime is None else int(optime[0])
self._cluster = Cluster(initialize, leader, self.last_leader_operation,
members, failover)
def _load_cluster(self):
if self.exhibitor and self.exhibitor.poll():
self.client.set_hosts(self.exhibitor.zookeeper_hosts)
if self.fetch_cluster:
try:
self.client.retry(self._inner_load_cluster)
except:
logger.exception('get_cluster')
self.session_listener(KazooState.LOST)
raise ZooKeeperError('ZooKeeper in not responding properly')
def _create(self, path, value, **kwargs):
try:
self.client.retry(self.client.create, path, value.encode('utf-8'),
**kwargs)
return True
except:
return False
def attempt_to_acquire_leader(self):
ret = self._create(self.leader_path,
self._name,
makepath=True,
ephemeral=True)
ret or logger.info('Could not take out TTL lock')
return ret
def set_failover_value(self, value, index=None):
try:
self.client.retry(self.client.set,
self.failover_path,
value.encode('utf-8'),
version=index or -1)
return True
except NoNodeError:
return value == '' or (not index
and self._create(self.failover_path, value))
except:
logging.exception('set_failover_value')
return False
def initialize(self, create_new=True, sysid=""):
return self._create(self.initialize_path, sysid, makepath=True) if create_new \
else self.client.retry(self.client.set, self.initialize_path, sysid.encode("utf-8"))
def touch_member(self, data, ttl=None):
cluster = self.cluster
me = cluster and ([m for m in cluster.members if m.name == self._name]
or [None])[0]
path = self.member_path
data = data.encode('utf-8')
create = not me
if me and self.client.client_id is not None and me.session != self.client.client_id[
0]:
try:
self.client.retry(self.client.delete, path)
except NoNodeError:
pass
except:
return False
create = True
if not create and data == self._my_member_data:
return True
try:
if create:
self.client.retry(self.client.create,
path,
data,
makepath=True,
ephemeral=True)
else:
self.client.retry(self.client.set, path, data)
self._my_member_data = data
return True
except NodeExistsError:
try:
self.client.retry(self.client.set, path, data)
self._my_member_data = data
return True
except:
logger.exception('touch_member')
except:
logger.exception('touch_member')
return False
def take_leader(self):
return self.attempt_to_acquire_leader()
def write_leader_optime(self, last_operation):
last_operation = last_operation.encode('utf-8')
if last_operation != self.last_leader_operation:
self.last_leader_operation = last_operation
path = self.leader_optime_path
try:
self.client.retry(self.client.set, path, last_operation)
except NoNodeError:
try:
self.client.retry(self.client.create,
path,
last_operation,
makepath=True)
except:
logger.exception('Failed to create %s', path)
except:
logger.exception('Failed to update %s', path)
def update_leader(self):
return True
def delete_leader(self):
self.client.restart()
self._my_member_data = None
return True
def _cancel_initialization(self):
node = self.get_node(self.initialize_path)
if node:
self.client.delete(self.initialize_path, version=node[1].version)
def cancel_initialization(self):
try:
self.client.retry(self._cancel_initialization)
except:
logger.exception("Unable to delete initialize key")
def watch(self, timeout):
if super(ZooKeeper, self).watch(timeout):
self.fetch_cluster = True
return self.fetch_cluster
class ZooKeeperJobStore(BaseJobStore):
"""
Stores jobs in a ZooKeeper tree. Any leftover keyword arguments are directly passed to
kazoo's `KazooClient
<http://kazoo.readthedocs.io/en/latest/api/client.html>`_.
Plugin alias: ``zookeeper``
:param str path: path to store jobs in
:param client: a :class:`~kazoo.client.KazooClient` instance to use instead of
providing connection arguments
:param int pickle_protocol: pickle protocol level to use (for serialization), defaults to the
highest available
"""
def __init__(self,
path='/apscheduler',
client=None,
close_connection_on_exit=False,
pickle_protocol=pickle.HIGHEST_PROTOCOL,
**connect_args):
super().__init__()
self.pickle_protocol = pickle_protocol
self.close_connection_on_exit = close_connection_on_exit
if not path:
raise ValueError('The "path" parameter must not be empty')
self.path = path
if client:
self.client = maybe_ref(client)
else:
self.client = KazooClient(**connect_args)
self._ensured_path = False
def _ensure_paths(self):
if not self._ensured_path:
self.client.ensure_path(self.path)
self._ensured_path = True
def start(self, scheduler, alias):
super().start(scheduler, alias)
if not self.client.connected:
self.client.start()
def lookup_job(self, job_id):
self._ensure_paths()
node_path = os.path.join(self.path, job_id)
try:
content, _ = self.client.get(node_path)
doc = pickle.loads(content)
job = self._reconstitute_job(doc['job_state'])
return job
except BaseException:
return None
def get_due_jobs(self, now):
timestamp = datetime_to_utc_timestamp(now)
jobs = [
job_def['job'] for job_def in self._get_jobs()
if job_def['next_run_time'] is not None
and job_def['next_run_time'] <= timestamp
]
return jobs
def get_next_run_time(self):
next_runs = [
job_def['next_run_time'] for job_def in self._get_jobs()
if job_def['next_run_time'] is not None
]
return utc_timestamp_to_datetime(
min(next_runs)) if len(next_runs) > 0 else None
def get_all_jobs(self):
jobs = [job_def['job'] for job_def in self._get_jobs()]
self._fix_paused_jobs_sorting(jobs)
return jobs
def add_job(self, job):
self._ensure_paths()
node_path = os.path.join(self.path, str(job.id))
value = {
'next_run_time': datetime_to_utc_timestamp(job.next_run_time),
'job_state': job.__getstate__()
}
data = pickle.dumps(value, self.pickle_protocol)
try:
self.client.create(node_path, value=data)
except NodeExistsError:
raise ConflictingIdError(job.id)
def update_job(self, job):
self._ensure_paths()
node_path = os.path.join(self.path, str(job.id))
changes = {
'next_run_time': datetime_to_utc_timestamp(job.next_run_time),
'job_state': job.__getstate__()
}
data = pickle.dumps(changes, self.pickle_protocol)
try:
self.client.set(node_path, value=data)
except NoNodeError:
raise JobLookupError(job.id)
def remove_job(self, job_id):
self._ensure_paths()
node_path = os.path.join(self.path, str(job_id))
try:
self.client.delete(node_path)
except NoNodeError:
raise JobLookupError(job_id)
def remove_all_jobs(self):
try:
self.client.delete(self.path, recursive=True)
except NoNodeError:
pass
self._ensured_path = False
def shutdown(self):
if self.close_connection_on_exit:
self.client.stop()
self.client.close()
def _reconstitute_job(self, job_state):
job_state = job_state
job = Job.__new__(Job)
job.__setstate__(job_state)
job._scheduler = self._scheduler
job._jobstore_alias = self._alias
return job
def _get_jobs(self):
self._ensure_paths()
jobs = []
failed_job_ids = []
all_ids = self.client.get_children(self.path)
for node_name in all_ids:
try:
node_path = os.path.join(self.path, node_name)
content, _ = self.client.get(node_path)
doc = pickle.loads(content)
job_def = {
'job_id':
node_name,
'next_run_time':
doc['next_run_time'] if doc['next_run_time'] else None,
'job_state':
doc['job_state'],
'job':
self._reconstitute_job(doc['job_state']),
'creation_time':
_.ctime
}
jobs.append(job_def)
except BaseException:
self._logger.exception(
'Unable to restore job "%s" -- removing it' % node_name)
failed_job_ids.append(node_name)
# Remove all the jobs we failed to restore
if failed_job_ids:
for failed_id in failed_job_ids:
self.remove_job(failed_id)
paused_sort_key = datetime(9999, 12, 31, tzinfo=utc)
return sorted(jobs,
key=lambda job_def:
(job_def['job'].next_run_time or paused_sort_key,
job_def['creation_time']))
def __repr__(self):
self._logger.exception('<%s (client=%s)>' %
(self.__class__.__name__, self.client))
return '<%s (client=%s)>' % (self.__class__.__name__, self.client)
class AnalyticsDiscovery(gevent.Greenlet):
def _sandesh_connection_info_update(self, status, message):
new_conn_state = getattr(ConnectionStatus, status)
ConnectionState.update(conn_type = ConnectionType.ZOOKEEPER,
name = self._svc_name, status = new_conn_state,
message = message,
server_addrs = self._zk_server.split(','))
if (self._conn_state and self._conn_state != ConnectionStatus.DOWN and
new_conn_state == ConnectionStatus.DOWN):
msg = 'Connection to Zookeeper down: %s' %(message)
self._logger.error(msg)
if (self._conn_state and self._conn_state != new_conn_state and
new_conn_state == ConnectionStatus.UP):
msg = 'Connection to Zookeeper ESTABLISHED'
self._logger.error(msg)
self._conn_state = new_conn_state
#import pdb; pdb.set_trace()
# end _sandesh_connection_info_update
def _zk_listen(self, state):
self._logger.error("Analytics Discovery listen %s" % str(state))
if state == KazooState.CONNECTED:
if self._conn_state != ConnectionStatus.UP:
self._sandesh_connection_info_update(status='UP', message='')
self._logger.error("Analytics Discovery to publish %s" % str(self._pubinfo))
self._reconnect = True
else:
self._logger.error("Analytics Discovery already connected")
else:
self._logger.error("Analytics Discovery NOT connected")
if self._conn_state == ConnectionStatus.UP:
self._sandesh_connection_info_update(status='DOWN', message='')
def _zk_datawatch(self, watcher, child, data, stat, event="unknown"):
self._logger.error(\
"Analytics Discovery %s ChildData : child %s, data %s, event %s" % \
(watcher, child, data, event))
if data:
data_dict = json.loads(data)
self._wchildren[watcher][child] = OrderedDict(sorted(data_dict.items()))
else:
if child in self._wchildren[watcher]:
del self._wchildren[watcher][child]
if self._watchers[watcher]:
self._pendingcb.add(watcher)
def _zk_watcher(self, watcher, children):
self._logger.error("Analytics Discovery Children %s" % children)
self._reconnect = True
def __init__(self, logger, zkservers, svc_name, inst,
watchers={}, zpostfix="", freq=10):
gevent.Greenlet.__init__(self)
self._svc_name = svc_name
self._inst = inst
self._zk_server = zkservers
# initialize logging and other stuff
if logger is None:
logging.basicConfig()
self._logger = logging
else:
self._logger = logger
self._conn_state = None
self._sandesh_connection_info_update(status='INIT', message='')
self._zk = KazooClient(hosts=zkservers)
self._pubinfo = None
self._watchers = watchers
self._wchildren = {}
self._pendingcb = set()
self._zpostfix = zpostfix
self._basepath = "/analytics-discovery-" + self._zpostfix
self._reconnect = None
self._freq = freq
def publish(self, pubinfo):
self._pubinfo = pubinfo
#import pdb; pdb.set_trace()
if self._conn_state == ConnectionStatus.UP:
try:
self._logger.error("ensure %s" % (self._basepath + "/" + self._svc_name))
self._logger.error("zk state %s (%s)" % (self._zk.state, self._zk.client_state))
self._zk.ensure_path(self._basepath + "/" + self._svc_name)
self._logger.error("check for %s/%s/%s" % \
(self._basepath, self._svc_name, self._inst))
if pubinfo is not None:
if self._zk.exists("%s/%s/%s" % \
(self._basepath, self._svc_name, self._inst)):
self._zk.set("%s/%s/%s" % \
(self._basepath, self._svc_name, self._inst),
self._pubinfo)
else:
self._zk.create("%s/%s/%s" % \
(self._basepath, self._svc_name, self._inst),
self._pubinfo, ephemeral=True)
else:
if self._zk.exists("%s/%s/%s" % \
(self._basepath, self._svc_name, self._inst)):
self._logger.error("withdrawing published info!")
self._zk.delete("%s/%s/%s" % \
(self._basepath, self._svc_name, self._inst))
except Exception as ex:
template = "Exception {0} in AnalyticsDiscovery publish. Args:\n{1!r}"
messag = template.format(type(ex).__name__, ex.args)
self._logger.error("%s : traceback %s for %s info %s" % \
(messag, traceback.format_exc(), self._svc_name, str(self._pubinfo)))
self._sandesh_connection_info_update(status='DOWN', message='')
self._reconnect = True
else:
self._logger.error("Analytics Discovery cannot publish while down")
def _run(self):
while True:
try:
self._zk.start()
break
except gevent.event.Timeout as e:
# Update connection info
self._sandesh_connection_info_update(status='DOWN',
message=str(e))
gevent.sleep(1)
# Zookeeper is also throwing exception due to delay in master election
except Exception as e:
# Update connection info
self._sandesh_connection_info_update(status='DOWN',
message=str(e))
gevent.sleep(1)
try:
# Update connection info
self._sandesh_connection_info_update(status='UP', message='')
self._reconnect = False
# Done connecting to ZooKeeper
self._zk.add_listener(self._zk_listen)
for wk in self._watchers.keys():
self._zk.ensure_path(self._basepath + "/" + wk)
self._wchildren[wk] = {}
self._zk.ChildrenWatch(self._basepath + "/" + wk,
partial(self._zk_watcher, wk))
# Trigger the initial publish
self._reconnect = True
while True:
try:
if not self._reconnect:
pending_list = list(self._pendingcb)
self._pendingcb = set()
for wk in pending_list:
if self._watchers[wk]:
self._watchers[wk](\
sorted(self._wchildren[wk].values()))
# If a reconnect happens during processing, don't lose it
while self._reconnect:
self._logger.error("Analytics Discovery %s reconnect" \
% self._svc_name)
self._reconnect = False
self._pendingcb = set()
self.publish(self._pubinfo)
for wk in self._watchers.keys():
self._zk.ensure_path(self._basepath + "/" + wk)
children = self._zk.get_children(self._basepath + "/" + wk)
old_children = set(self._wchildren[wk].keys())
new_children = set(children)
# Remove contents for the children who are gone
# (DO NOT remove the watch)
for elem in old_children - new_children:
del self._wchildren[wk][elem]
# Overwrite existing children, or create new ones
for elem in new_children:
# Create a watch for new children
if elem not in self._wchildren[wk]:
self._zk.DataWatch(self._basepath + "/" + \
wk + "/" + elem,
partial(self._zk_datawatch, wk, elem))
data_str, _ = self._zk.get(\
self._basepath + "/" + wk + "/" + elem)
data_dict = json.loads(data_str)
self._wchildren[wk][elem] = \
OrderedDict(sorted(data_dict.items()))
self._logger.error(\
"Analytics Discovery %s ChildData : child %s, data %s, event %s" % \
(wk, elem, self._wchildren[wk][elem], "GET"))
if self._watchers[wk]:
self._watchers[wk](sorted(self._wchildren[wk].values()))
gevent.sleep(self._freq)
except gevent.GreenletExit:
self._logger.error("Exiting AnalyticsDiscovery for %s" % \
self._svc_name)
self._zk.stop()
break
except Exception as ex:
template = "Exception {0} in AnalyticsDiscovery reconnect. Args:\n{1!r}"
messag = template.format(type(ex).__name__, ex.args)
self._logger.error("%s : traceback %s for %s info %s" % \
(messag, traceback.format_exc(), self._svc_name, str(self._pubinfo)))
self._reconnect = True
except Exception as ex:
template = "Exception {0} in AnalyticsDiscovery run. Args:\n{1!r}"
messag = template.format(type(ex).__name__, ex.args)
self._logger.error("%s : traceback %s for %s info %s" % \
(messag, traceback.format_exc(), self._svc_name, str(self._pubinfo)))
raise SystemExit
from kazoo.client import KazooClient
zk = KazooClient('zk:2181')
zk.start()
zk.delete('/lymph', recursive=True)
zk.stop()
class SimpleSwitch13(app_manager.RyuApp):
OFP_VERSIONS = [ofproto_v1_3.OFP_VERSION]
def __init__(self, *args, **kwargs):
super(SimpleSwitch13, self).__init__(*args, **kwargs)
self.mac_to_port = {}
self.zkConf = {
'root': '/multicontroller',
'topo': '/topology',
'swstat': '/swstat',
'counter': '/counter'
}
self.zk = KazooClient('127.0.0.1:2181')
self.zk.start()
self.ip = '202.201.3.51'
self.sws = {}
self.gid = random.randint(0, 10000)
self.dps = {}
self.links = []
self.interval = 5
self.role = OFPCR_ROLE_EQUAL
self.topoThread = hub.spawn(self._topoThread)
self.linkThread = hub.spawn(self._linkDiscover)
self.clearLinkThread = hub.spawn(self._cleanLinks)
self.clearLinkThread = hub.spawn(self._cleanSwitches)
def _cleanSwitches(self):
while True:
self.sws = {k: self.sws[k] for k in self.sws if self.sws[k]}
hub.sleep(self.interval)
def _topoThread(self):
linkNode = self.zkConf['root'] + self.zkConf['topo'] + self.ip
if self.zk.exists(linkNode):
self.zk.set(linkNode, json.dumps(self.links))
else:
self.zk.create(linkNode, json.dumps(self.links))
hub.sleep(self.interval)
def _linkDiscover(self):
while True:
for dpid in self.dps:
self.sendSlldp(dpid)
hub.sleep(self.interval)
def sendSlldp(self, dpid):
dp = self.dps.get(dpid)
if dp is None:
return
actions = [dp.ofproto_parser.OFPActionOutput(dp.ofproto.OFPP_FLOOD)]
pkt = packet.Packet()
pkt.add_protocol(
ethernet.ethernet(ethertype=ETH_TYPE_SLLDP,
dst=SLLDP_MAC_DST,
src=SLLDP_MAC_SRC))
pkt.add_protocol(slldp(dp.id))
pkt.serialize()
slldpPacket = pkt.data
out = dp.ofproto_parser.OFPPacketOut(
datapath=dp,
in_port=dp.ofproto.OFPP_CONTROLLER,
buffer_id=dp.ofproto.OFP_NO_BUFFER,
actions=actions,
data=slldpPacket)
dp.send_msg(out)
def getLinks(self):
topoNode = self.zkConf['root'] + self.zkConf['topo']
ips = self.zk.get_children(topoNode)
res = []
for ip in ips:
links = self.zk.get(topoNode + ip)[0]
for link in links:
res.append(link)
return res
@set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
def _packet_in_handler(self, ev):
msg = ev.msg
pkt = packet.Packet(msg.data)
eth = pkt.get_protocols(ethernet.ethernet)[0]
dst = eth.dst
# SLLDP packet
if dst == SLLDP_MAC_DST:
self.handleSlldp(ev)
return
# process packet_in message in subclass
self.packet_in_process(ev)
def handleSlldp(self, ev):
msg = ev.msg
datapath = msg.datapath
dpid = datapath.id
inPort = msg.match['in_port']
pkt = packet.Packet(msg.data)
slldpBuff = pkt.get_protocols(ethernet.ethernet)[2]
dpidSrc, _ = slldp.parser(slldpBuff)
self.links.append({
'srcdpid': dpidSrc,
'dst': {
'dpid': dpid,
'port': inPort
},
'time': time.time()
})
def _cleanLinks(self):
while True:
now = time.time()
self.links = [
l for l in self.links if now - l['time'] < self.interval
]
hub.sleep(self.interval)
@abc.abstractmethod
def packet_in_process(self, ev):
pass
@set_ev_cls(event.EventSwitchEnter)
def switch_enter(self, ev):
dpid = ev.datapath.id
self.sws[dpid] = True
self.dps[dpid] = ev.datapath
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
self.zk.ensure_path(dpNode)
if self.election(dpid):
self.role = OFPCR_ROLE_MASTER
else:
self.role = OFPCR_ROLE_SLAVE
self.countUp(dpid)
self.roleRequest(dpid, self.role)
mflag = dpNode + '/' + 'master'
DataWatch(self.zk, mflag, self.masterWatcher)
def masterWatcher(self, data, stat, ev):
if ev and ev.type == 'DELETED':
_, _, dpid, _ = ev.path.split('/')
dpid = str_to_dpid(dpid)
if self.sws.get(dpid):
if self.election(dpid):
self.role = OFPCR_ROLE_MASTER
self.roleRequest(dpid, self.role)
return self.sws.get(dpid)
def election(self, dpid):
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
mflag = dpNode + '/' + 'master'
while not self.zk.exists(mflag):
mlock = self.zk.Lock(dpNode + '/' + 'mlock', self.ip)
with mlock:
if not self.zk.exists(mflag):
self.zk.create(mflag, self.ip, ephemeral=True)
if self.zk.exists(mflag):
if self.zk.get(mflag) == self.ip:
return True
else:
return False
else:
time.sleep(random.randint(0, 100) / 500.0)
return False
def roleRequest(self, dp, role):
msg = dp.ofproto_parser.OFPRoleRequest(dp, role, self.gid)
dp.send_msg(msg)
def getCount(self, dpid):
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
countNode = dpNode + self.zkConf['counter']
counters = self.zk.get_children(countNode)
return len(counters)
def countUp(self, dpid):
countNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid) + self.zkConf['counter']
self.zk.ensure_path(countNode)
self.zk.create(countNode + uuid4().hex, 'alive', ephemeral=True)
@set_ev_cls(event.EventSwitchLeave)
def switch_levave(self, ev):
dpid = ev.datapath
count = self.getCount(dpid)
self.sws[dpid] = False
if count == 0:
dpNode = self.zkConf['root'] + self.zkConf['swstat'] \
+ '/' + dpid_to_str(dpid)
self.zk.delete(dpNode, recursive=True)
class ArcusZooKeeper:
"""
ZooKeeper helper for Arcus
"""
def __init__(self, hostports, timeout):
self.hostports = hostports
self.timeout = timeout
self.zk = KazooClient(hosts=hostports, read_only=False)
def start(self):
self.zk.start()
def stop(self):
self.zk.stop()
def init_structure(self):
if self.zk.exists("/arcus"):
print "init_arcus_structure: fail (/arcus exists)"
return False
tx = self.zk.transaction()
tx.create("/arcus", b"")
tx.create("/arcus/cache_list", b"")
tx.create("/arcus/client_list", b"")
tx.create("/arcus/cache_server_mapping", b"")
results = tx.commit()
if len(results) > 0:
print results
return False
print "init_structure: success"
return True
def drop_structure(self):
self.zk.delete("/arcus", recursive=True)
print "delete_structure: success"
def get_structure(self):
return self.zk.get_children("/arcus")
def get_mapping_for_service(self, service_code):
result = []
cache_list = "/arcus/cache_list/%s" % service_code
mapping = "/arcus/cache_server_mapping"
all = self.zk.get_children(mapping)
for ipport in all:
codes = self.zk.get_children("%s/%s" % (mapping, ipport))
if len(codes) > 0:
if codes[0] == service_code:
result.append("%s/%s" % (mapping, ipport))
return result
def get_config_for_service(self, service_code):
cache_list = "/arcus/cache_list/%s" % service_code
data, stat = self.zk.get(cache_list)
return json.loads(data), data, stat
def update_service_code(self, cluster):
cache_list = "/arcus/cache_list/%s" % cluster["serviceCode"]
client_list = "/arcus/client_list/%s" % cluster["serviceCode"]
mapping = "/arcus/cache_server_mapping"
try:
delete_list = self.get_mapping_for_service(cluster["serviceCode"])
# 0. Create a transaction
tx = self.zk.transaction()
# 1. Cache list
if self.zk.exists(cache_list):
c1, c2, c3 = self.get_config_for_service(cluster["serviceCode"])
cluster["created"] = c1.get("created")
cluster["modified"] = str(datetime.datetime.now())
tx.set_data(cache_list, json.dumps(cluster))
else:
cluster["created"] = str(datetime.datetime.now())
tx.create("/arcus/cache_list/%s" % cluster["serviceCode"], json.dumps(cluster))
# 2. Client list
if self.zk.exists(client_list):
pass
else:
tx.create("/arcus/client_list/%s" % cluster["serviceCode"], b"")
# 3. Mapping
for each in delete_list:
tx.delete("%s/%s" % (each, cluster["serviceCode"]))
tx.delete(each)
for server in cluster["servers"]:
global_config = cluster.get("config", {})
per_node_config = server.get("config", {})
config = dict(global_config.items() + per_node_config.items())
if len(config) == 0:
print "update_service_code: config not found for {0}".format(server)
continue
map_ip = "/arcus/cache_server_mapping/%s:%s" % (server["ip"], config["port"])
map_code = "%s/%s" % (map_ip, cluster["serviceCode"])
tx.create(map_ip, json.dumps(config))
tx.create(map_code, b"")
# 4. Commit
results = tx.commit()
print results
except Exception as e:
traceback.print_exc()
def delete_service_code(self, cluster):
cache_list = "/arcus/cache_list/%s" % cluster["serviceCode"]
client_list = "/arcus/client_list/%s" % cluster["serviceCode"]
mapping = "/arcus/cache_server_mapping"
try:
delete_list = self.get_mapping_for_service(cluster["serviceCode"])
# 0. Create a transaction
tx = self.zk.transaction()
# 1. Cache list
tx.delete("/arcus/cache_list/%s" % cluster["serviceCode"])
# 2. Client list
tx.delete("/arcus/client_list/%s" % cluster["serviceCode"])
# 3. Mapping
for each in delete_list:
tx.delete("%s/%s" % (each, cluster["serviceCode"]))
tx.delete(each)
# 4. Commit
results = tx.commit()
print results
except Exception as e:
traceback.print_exc()
def list_all_service_code(self):
result = []
cache_list = "/arcus/cache_list"
try:
list = self.zk.get_children(cache_list)
for each in list:
result.append(self.list_service_code(each))
return result
except Exception as e:
traceback.print_exc()
def list_service_code(self, service_code):
result = {}
cache_list = "/arcus/cache_list/%s" % service_code
client_list = "/arcus/client_list/%s" % service_code
mapping = "/arcus/cache_server_mapping"
try:
data, stat = self.zk.get(cache_list)
static_list = self.get_mapping_for_service(service_code)
current_list = self.zk.get_children(cache_list)
# sort the lists
static_list.sort()
current_list.sort()
# get clusterConfig
cluster = json.loads(data)
# get clusterStatus
static_set = set([each.split("/")[-1] for each in static_list])
current_set = set([each.split("-")[0] for each in current_list])
offline = static_set - current_set
online = static_set - offline
undefined = current_set - static_set
result["serviceCode"] = service_code
result["config"] = cluster
result["online"] = list(online)
result["offline"] = list(offline)
result["undefined"] = list(undefined)
result["created"] = cluster.get("created")
result["modified"] = cluster.get("modified")
return result
except Exception as e:
traceback.print_exc()
class ZooKeeperClient:
"""
Handles basic ZooKeeper events
"""
def __init__(self, zk_hosts, log_name="ZooKeeperClient", prefix=""):
"""
:param zk_hosts: List of ZooKeepers hosts
:param log_name: Name of the logger to use
:param prefix: Prefix to all paths
"""
self._zk = KazooClient(zk_hosts)
self._zk.add_listener(self.__conn_listener)
self.__prefix = prefix
# Session state
self.__connected = False
self.__online = False
self.__stop = False
# Notification queue
self._logger = logging.getLogger(log_name)
self._queue = ThreadPool(1, 1, logname=log_name)
@property
def prefix(self):
"""
Prefix to all ZooKeeper nodes
"""
return self.__prefix
@property
def connected(self):
"""
ZooKeeper client state: connected to the quorum
"""
return self.__connected
@property
def online(self):
"""
ZooKeeper client state: connected & online (session active)
"""
return self.__online
@property
def stopped(self):
"""
ZooKeeper client status (stop requested)
"""
return self.__stop
def __conn_listener(self, state):
"""
Connection event listener
:param state: The new connection state
"""
if state == KazooState.CONNECTED:
self.__online = True
if not self.__connected:
self.__connected = True
self._logger.info("Connected to ZooKeeper")
self._queue.enqueue(self.on_first_connection)
else:
self._logger.warning("Re-connected to ZooKeeper")
self._queue.enqueue(self.on_client_reconnection)
elif state == KazooState.SUSPENDED:
self._logger.warning("Connection suspended")
self.__online = False
elif state == KazooState.LOST:
self.__online = False
self.__connected = False
if self.__stop:
self._logger.info("Disconnected from ZooKeeper (requested)")
else:
self._logger.warning("Connection lost")
def start(self):
"""
Starts the connection
"""
self.__stop = False
self._queue.start()
self._zk.start()
def stop(self):
"""
Stops the connection
"""
self.__stop = True
self._queue.stop()
self._zk.stop()
@staticmethod
def on_first_connection():
"""
Called when the client is connected for the first time
"""
pass
@staticmethod
def on_client_reconnection():
"""
Called when the client is reconnected to the server
"""
pass
def __path(self, path):
"""
Adds the prefix to the given path
:param path: Z-Path
:return: Prefixed Z-Path
"""
if path.startswith(self.__prefix):
return path
return "{}{}".format(self.__prefix, path)
def create(self, path, data, ephemeral=False, sequence=False):
"""
Creates a ZooKeeper node
:param path: Z-Path
:param data: Node Content
:param ephemeral: Ephemeral flag
:param sequence: Sequential flag
"""
return self._zk.create(self.__path(path),
data,
ephemeral=ephemeral,
sequence=sequence)
def ensure_path(self, path):
"""
Ensures that a path exists, creates it if necessary
:param path: Z-Path
"""
return self._zk.ensure_path(self.__path(path))
def get(self, path, watch=None):
"""
Gets the content of a ZooKeeper node
:param path: Z-Path
:param watch: Watch method
"""
return self._zk.get(self.__path(path), watch=watch)
def get_children(self, path, watch=None):
"""
Gets the list of children of a node
:param path: Z-Path
:param watch: Watch method
"""
return self._zk.get_children(self.__path(path), watch=watch)
def set(self, path, data):
"""
Sets the content of a ZooKeeper node
:param path: Z-Path
:param data: New content
"""
return self._zk.set(self.__path(path), data)
def delete(self, path):
"""
Deletes a node
:param path: Z-Path
"""
return self._zk.delete(self.__path(path))
class ZooKeeper(AbstractDCS):
def __init__(self, config):
super(ZooKeeper, self).__init__(config)
hosts = config.get('hosts', [])
if isinstance(hosts, list):
hosts = ','.join(hosts)
self._client = KazooClient(hosts, handler=PatroniSequentialThreadingHandler(config['retry_timeout']),
timeout=config['ttl'], connection_retry={'max_delay': 1, 'max_tries': -1},
command_retry={'deadline': config['retry_timeout'], 'max_delay': 1, 'max_tries': -1})
self._client.add_listener(self.session_listener)
self._my_member_data = None
self._fetch_cluster = True
self._last_leader_operation = 0
self._orig_kazoo_connect = self._client._connection._connect
self._client._connection._connect = self._kazoo_connect
self._client.start()
def _kazoo_connect(self, host, port):
"""Kazoo is using Ping's to determine health of connection to zookeeper. If there is no
response on Ping after Ping interval (1/2 from read_timeout) it will consider current
connection dead and try to connect to another node. Without this "magic" it was taking
up to 2/3 from session timeout (ttl) to figure out that connection was dead and we had
only small time for reconnect and retry.
This method is needed to return different value of read_timeout, which is not calculated
from negotiated session timeout but from value of `loop_wait`. And it is 2 sec smaller
than loop_wait, because we can spend up to 2 seconds when calling `touch_member()` and
`write_leader_optime()` methods, which also may hang..."""
ret = self._orig_kazoo_connect(host, port)
return max(self.loop_wait - 2, 2)*1000, ret[1]
def session_listener(self, state):
if state in [KazooState.SUSPENDED, KazooState.LOST]:
self.cluster_watcher(None)
def cluster_watcher(self, event):
self._fetch_cluster = True
self.event.set()
def reload_config(self, config):
self.set_retry_timeout(config['retry_timeout'])
loop_wait = config['loop_wait']
loop_wait_changed = self._loop_wait != loop_wait
self._loop_wait = loop_wait
self._client.handler.set_connect_timeout(loop_wait)
# We need to reestablish connection to zookeeper if we want to change
# read_timeout (and Ping interval respectively), because read_timeout
# is calculated in `_kazoo_connect` method. If we are changing ttl at
# the same time, set_ttl method will reestablish connection and return
# `!True`, otherwise we will close existing connection and let kazoo
# open the new one.
if not self.set_ttl(int(config['ttl'] * 1000)) and loop_wait_changed:
self._client._connection._socket.close()
def set_ttl(self, ttl):
"""It is not possible to change ttl (session_timeout) in zookeeper without
destroying old session and creating the new one. This method returns `!True`
if session_timeout has been changed (`restart()` has been called)."""
if self._client._session_timeout != ttl:
self._client._session_timeout = ttl
self._client.restart()
return True
def set_retry_timeout(self, retry_timeout):
self._client._retry.deadline = retry_timeout
def get_node(self, key, watch=None):
try:
ret = self._client.get(key, watch)
return (ret[0].decode('utf-8'), ret[1])
except NoNodeError:
return None
@staticmethod
def member(name, value, znode):
return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
def get_children(self, key, watch=None):
try:
return self._client.get_children(key, watch)
except NoNodeError:
return []
def load_members(self):
members = []
for member in self.get_children(self.members_path, self.cluster_watcher):
data = self.get_node(self.members_path + member)
if data is not None:
members.append(self.member(member, *data))
return members
def _inner_load_cluster(self):
self._fetch_cluster = False
self.event.clear()
nodes = set(self.get_children(self.client_path(''), self.cluster_watcher))
if not nodes:
self._fetch_cluster = True
# get initialize flag
initialize = (self.get_node(self.initialize_path) or [None])[0] if self._INITIALIZE in nodes else None
# get global dynamic configuration
config = self.get_node(self.config_path, watch=self.cluster_watcher) if self._CONFIG in nodes else None
config = config and ClusterConfig.from_node(config[1].version, config[0], config[1].mzxid)
# get list of members
members = self.load_members() if self._MEMBERS[:-1] in nodes else []
# get leader
leader = self.get_node(self.leader_path) if self._LEADER in nodes else None
if leader:
client_id = self._client.client_id
if leader[0] == self._name and client_id is not None and client_id[0] != leader[1].ephemeralOwner:
logger.info('I am leader but not owner of the session. Removing leader node')
self._client.delete(self.leader_path)
leader = None
if leader:
member = Member(-1, leader[0], None, {})
member = ([m for m in members if m.name == leader[0]] or [member])[0]
leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
self._fetch_cluster = member.index == -1
# failover key
failover = self.get_node(self.failover_path, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
failover = failover and Failover.from_node(failover[1].version, failover[0])
# get last leader operation
optime = self.get_node(self.leader_optime_path) if self._OPTIME in nodes and self._fetch_cluster else None
self._last_leader_operation = 0 if optime is None else int(optime[0])
self._cluster = Cluster(initialize, config, leader, self._last_leader_operation, members, failover)
def _load_cluster(self):
if self._fetch_cluster or self._cluster is None:
try:
self._client.retry(self._inner_load_cluster)
except Exception:
logger.exception('get_cluster')
self.cluster_watcher(None)
raise ZooKeeperError('ZooKeeper in not responding properly')
def _create(self, path, value, **kwargs):
try:
self._client.retry(self._client.create, path, value.encode('utf-8'), **kwargs)
return True
except:
return False
def attempt_to_acquire_leader(self, permanent=False):
ret = self._create(self.leader_path, self._name, makepath=True, ephemeral=not permanent)
if not ret:
logger.info('Could not take out TTL lock')
return ret
def set_failover_value(self, value, index=None):
try:
self._client.retry(self._client.set, self.failover_path, value.encode('utf-8'), version=index or -1)
return True
except NoNodeError:
return value == '' or (index is None and self._create(self.failover_path, value))
except:
logging.exception('set_failover_value')
return False
def set_config_value(self, value, index=None):
try:
self._client.retry(self._client.set, self.config_path, value.encode('utf-8'), version=index or -1)
return True
except NoNodeError:
return index is None and self._create(self.config_path, value)
except Exception:
logging.exception('set_config_value')
return False
def initialize(self, create_new=True, sysid=""):
return self._create(self.initialize_path, sysid, makepath=True) if create_new \
else self._client.retry(self._client.set, self.initialize_path, sysid.encode("utf-8"))
def touch_member(self, data, ttl=None, permanent=False):
cluster = self.cluster
member = cluster and ([m for m in cluster.members if m.name == self._name] or [None])[0]
data = data.encode('utf-8')
if member and self._client.client_id is not None and member.session != self._client.client_id[0]:
try:
self._client.delete_async(self.member_path).get(timeout=1)
except NoNodeError:
pass
except:
return False
member = None
if member:
if data == self._my_member_data:
return True
else:
try:
self._client.create_async(self.member_path, data, makepath=True, ephemeral=not permanent).get(timeout=1)
self._my_member_data = data
return True
except Exception as e:
if not isinstance(e, NodeExistsError):
logger.exception('touch_member')
return False
try:
self._client.set_async(self.member_path, data).get(timeout=1)
self._my_member_data = data
return True
except:
logger.exception('touch_member')
return False
def take_leader(self):
return self.attempt_to_acquire_leader()
def write_leader_optime(self, last_operation):
last_operation = last_operation.encode('utf-8')
if last_operation != self._last_leader_operation:
try:
self._client.set_async(self.leader_optime_path, last_operation).get(timeout=1)
self._last_leader_operation = last_operation
except NoNodeError:
try:
self._client.create_async(self.leader_optime_path, last_operation, makepath=True).get(timeout=1)
self._last_leader_operation = last_operation
except:
logger.exception('Failed to create %s', self.leader_optime_path)
except:
logger.exception('Failed to update %s', self.leader_optime_path)
def update_leader(self):
return True
def delete_leader(self):
self._client.restart()
self._my_member_data = None
return True
def _cancel_initialization(self):
node = self.get_node(self.initialize_path)
if node:
self._client.delete(self.initialize_path, version=node[1].version)
def cancel_initialization(self):
try:
self._client.retry(self._cancel_initialization)
except:
logger.exception("Unable to delete initialize key")
def delete_cluster(self):
try:
return self._client.retry(self._client.delete, self.client_path(''), recursive=True)
except NoNodeError:
return True
def watch(self, timeout):
if super(ZooKeeper, self).watch(timeout):
self._fetch_cluster = True
return self._fetch_cluster
class ZK:
"""Opens a connection to a kafka zookeeper. "
"To be used in the 'with' statement."""
def __init__(self, cluster_config):
self.cluster_config = cluster_config
def __enter__(self):
kazooRetry = KazooRetry(max_tries=5, )
self.zk = KazooClient(
hosts=self.cluster_config.zookeeper,
read_only=True,
connection_retry=kazooRetry,
)
_log.debug(
"ZK: Creating new zookeeper connection: {zookeeper}".format(
zookeeper=self.cluster_config.zookeeper), )
self.zk.start()
return self
def __exit__(self, type, value, traceback):
self.zk.stop()
def get_children(self, path, watch=None):
"""Returns the children of the specified node."""
_log.debug("ZK: Getting children of {path}".format(path=path), )
return self.zk.get_children(path, watch)
def get(self, path, watch=None):
"""Returns the data of the specified node."""
_log.debug("ZK: Getting {path}".format(path=path), )
return self.zk.get(path, watch)
def set(self, path, value):
"""Sets and returns new data for the specified node."""
_log.debug("ZK: Setting {path} to {value}".format(path=path,
value=value))
return self.zk.set(path, value)
def get_json(self, path, watch=None):
"""Reads the data of the specified node and converts it to json."""
data, _ = self.get(path, watch)
return json.loads(data) if data else None
def get_broker_metadata(self, broker_id):
try:
broker_json, _ = self.get(
"/brokers/ids/{b_id}".format(b_id=broker_id))
except NoNodeError:
_log.error("broker '{b_id}' not found.".format(b_id=broker_id), )
raise
return json.loads(broker_json)
def get_brokers(self, names_only=False):
"""Get information on all the available brokers.
:rtype : dict of brokers
"""
broker_ids = self.get_children("/brokers/ids")
# Return broker-ids only
if names_only:
return {int(b_id): None for b_id in broker_ids}
return {
int(b_id): self.get_broker_metadata(b_id)
for b_id in broker_ids
}
def get_topic_config(self, topic):
"""Get configuration information for specified topic.
:rtype : dict of configuration
"""
try:
config_data = json.loads(
self.get("/config/topics/{topic}".format(topic=topic))[0])
except NoNodeError as e:
_log.error("topic {topic} not found.".format(topic=topic))
raise e
return config_data
def set_topic_config(self, topic, value):
"""Set configuration information for specified topic.
:rtype : dict of new configuration"""
try:
config_data = json.dumps(value)
# Change value
return_value = self.set(
"/config/topics/{topic}".format(topic=topic), config_data)
# Create change
self.create('/config/changes/config_change_', topic, sequence=True)
except NoNodeError as e:
_log.error("topic {topic} not found.".format(topic=topic))
raise e
return return_value
def get_topics(
self,
topic_name=None,
names_only=False,
fetch_partition_state=True,
):
"""Get information on all the available topics.
Topic-data format with fetch_partition_state as False :-
topic_data = {
'version': 1,
'partitions': {
<p_id>: {
replicas: <broker-ids>
}
}
}
Topic-data format with fetch_partition_state as True:-
topic_data = {
'version': 1,
'partitions': {
<p_id>:{
replicas: [<broker_id>, <broker_id>, ...],
isr: [<broker_id>, <broker_id>, ...],
controller_epoch: <val>,
leader_epoch: <val>,
version: 1,
leader: <broker-id>,
}
}
Note: By default we also fetch partition-state which results in
accessing the zookeeper twice. If just partition-replica information is
required fetch_partition_state should be set to False.
"""
topic_ids = [topic_name] if topic_name else self.get_children(
"/brokers/topics", )
if names_only:
return topic_ids
topics_data = {}
for topic_id in topic_ids:
try:
topic_data = json.loads(
self.get("/brokers/topics/{id}".format(id=topic_id))[0], )
except NoNodeError:
_log.error(
"topic '{topic}' not found.".format(topic=topic_id), )
return {}
# Prepare data for each partition
partitions_data = {}
for p_id, replicas in topic_data['partitions'].iteritems():
partitions_data[p_id] = {}
if fetch_partition_state:
# Fetch partition-state from zookeeper
partitions_data[p_id] = self._fetch_partition_state(
topic_id, p_id)
partitions_data[p_id]['replicas'] = replicas
topic_data['partitions'] = partitions_data
topics_data[topic_id] = topic_data
return topics_data
def get_consumer_groups(self, consumer_group_id=None, names_only=False):
"""Get information on all the available consumer-groups.
If names_only is False, only list of consumer-group ids are sent.
If names_only is True, Consumer group offset details are returned
for all consumer-groups or given consumer-group if given in dict
format as:-
{
'group-id':
{
'topic':
{
'partition': offset-value,
...
...
}
}
}
:rtype: dict of consumer-group offset details
"""
if consumer_group_id is None:
group_ids = self.get_children("/consumers")
else:
group_ids = [consumer_group_id]
# Return consumer-group-ids only
if names_only:
return {g_id: None for g_id in group_ids}
consumer_offsets = {}
for g_id in group_ids:
consumer_offsets[g_id] = self.get_group_offsets(g_id)
return consumer_offsets
def get_group_offsets(self, group, topic=None):
"""Fetch group offsets for given topic and partition otherwise all topics
and partitions otherwise.
{
'topic':
{
'partition': offset-value,
...
...
}
}
"""
group_offsets = {}
try:
all_topics = self.get_my_subscribed_topics(group)
except NoNodeError:
# No offset information of given consumer-group
_log.warning(
"No topics subscribed to consumer-group {group}.".format(
group=group, ), )
return group_offsets
if topic:
if topic in all_topics:
topics = [topic]
else:
_log.error(
"Topic {topic} not found in topic list {topics} for consumer"
"-group {consumer_group}.".format(
topic=topic,
topics=', '.join(topic for topic in all_topics),
consumer_group=group,
), )
return group_offsets
else:
topics = all_topics
for topic in topics:
group_offsets[topic] = {}
try:
partitions = self.get_my_subscribed_partitions(group, topic)
except NoNodeError:
_log.warning(
"No partition offsets found for topic {topic}. "
"Continuing to next one...".format(topic=topic), )
continue
# Fetch offsets for each partition
for partition in partitions:
path = "/consumers/{group_id}/offsets/{topic}/{partition}".format(
group_id=group,
topic=topic,
partition=partition,
)
try:
# Get current offset
offset_json, _ = self.get(path)
group_offsets[topic][partition] = json.loads(offset_json)
except NoNodeError:
_log.error("Path {path} not found".format(path=path))
raise
return group_offsets
def _fetch_partition_state(self, topic_id, partition_id):
"""Fetch partition-state for given topic-partition."""
state_path = "/brokers/topics/{topic_id}/partitions/{p_id}/state"
try:
partition_json, _ = self.get(
state_path.format(topic_id=topic_id, p_id=partition_id), )
return json.loads(partition_json)
except NoNodeError:
return {} # The partition has no data
def get_my_subscribed_topics(self, groupid):
"""Get the list of topics that a consumer is subscribed to
:param: groupid: The consumer group ID for the consumer
:returns list of kafka topics
:rtype: list
"""
path = "/consumers/{group_id}/offsets".format(group_id=groupid)
return self.get_children(path)
def get_my_subscribed_partitions(self, groupid, topic):
"""Get the list of partitions of a topic
that a consumer is subscribed to
:param: groupid: The consumer group ID for the consumer
:param: topic: The topic name
:returns list of partitions
:rtype: list
"""
path = "/consumers/{group_id}/offsets/{topic}".format(
group_id=groupid,
topic=topic,
)
return self.get_children(path)
def get_cluster_assignment(self):
"""Fetch the cluster layout in form of assignment from zookeeper"""
plan = self.get_cluster_plan()
assignment = {}
for elem in plan['partitions']:
assignment[(elem['topic'], elem['partition'])] = elem['replicas']
return assignment
def create(self,
path,
value='',
acl=None,
ephemeral=False,
sequence=False,
makepath=False):
"""Creates a Zookeeper node.
:param: path: The zookeeper node path
:param: value: Zookeeper node value
:param: acl: ACL list
:param: ephemeral: Boolean indicating where this node is tied to
this session.
:param: sequence: Boolean indicating whether path is suffixed
with a unique index.
:param: makepath: Whether the path should be created if it doesn't
exist.
"""
_log.debug("ZK: Creating node " + path)
return self.zk.create(path, value, acl, ephemeral, sequence, makepath)
def delete(self, path, recursive=False):
"""Deletes a Zookeeper node.
:param: path: The zookeeper node path
:param: recursive: Recursively delete node and all its children.
"""
_log.debug("ZK: Deleting node " + path)
return self.zk.delete(path, recursive=recursive)
def delete_topic_partitions(self, groupid, topic, partitions):
"""Delete the specified partitions within the topic that the consumer
is subscribed to.
:param: groupid: The consumer group ID for the consumer.
:param: topic: Kafka topic.
:param: partitions: List of partitions within the topic to be deleted.
:raises:
NoNodeError: if the consumer is not subscribed to the topic
ZookeeperError: if there is an error with Zookeeper
"""
for partition in partitions:
path = "/consumers/{groupid}/offsets/{topic}/{partition}".format(
groupid=groupid, topic=topic, partition=partition)
self.delete(path)
def delete_topic(self, groupid, topic):
path = "/consumers/{groupid}/offsets/{topic}".format(
groupid=groupid,
topic=topic,
)
self.delete(path, True)
def delete_group(self, groupid):
path = "/consumers/{groupid}".format(groupid=groupid, )
self.delete(path, True)
def execute_plan(self, plan):
"""Submit reassignment plan for execution."""
reassignment_path = '{admin}/{reassignment_node}'\
.format(admin=ADMIN_PATH, reassignment_node=REASSIGNMENT_NODE)
plan_json = json.dumps(plan)
base_plan = self.get_cluster_plan()
if not validate_plan(plan, base_plan):
_log.error('Given plan is invalid. ABORTING reassignment...')
return False
# Send proposed-plan to zookeeper
try:
_log.info('Sending plan to Zookeeper...')
self.create(reassignment_path, plan_json, makepath=True)
_log.info(
'Re-assign partitions node in Zookeeper updated successfully '
'with {plan}'.format(plan=plan), )
return True
except NodeExistsError:
_log.warning('Previous plan in progress. Exiting..')
in_progress_plan = json.loads(self.get(reassignment_path)[0])
in_progress_partitions = [
'{topic}-{p_id}'.format(
topic=p_data['topic'],
p_id=str(p_data['partition']),
) for p_data in in_progress_plan['partitions']
]
_log.warning(
'{count} partition(s) reassignment currently in progress:-'.
format(count=len(in_progress_partitions)), )
_log.warning(
'{partitions}. ABORTING reassignment...'.format(
partitions=', '.join(in_progress_partitions), ), )
return False
except Exception as e:
_log.error(
'Could not re-assign partitions {plan}. Error: {e}'.format(
plan=plan, e=e), )
return False
def get_cluster_plan(self):
"""Fetch cluster plan from zookeeper."""
_log.info('Fetching current cluster-topology from Zookeeper...')
cluster_layout = self.get_topics(fetch_partition_state=False)
# Re-format cluster-layout
partitions = [
{
'topic': topic_id,
'partition': int(p_id),
'replicas': partitions_data['replicas']
} for topic_id, topic_info in cluster_layout.iteritems()
for p_id, partitions_data in topic_info['partitions'].iteritems()
]
return {'version': 1, 'partitions': partitions}
def get_pending_plan(self):
"""Read the currently running plan on reassign_partitions node."""
reassignment_path = '{admin}/{reassignment_node}'\
.format(admin=ADMIN_PATH, reassignment_node=REASSIGNMENT_NODE)
try:
result = self.get(reassignment_path)
return json.loads(result[0])
except NoNodeError:
return {}
class BalancedConsumer(object):
"""
A self-balancing consumer for Kafka that uses ZooKeeper to communicate
with other balancing consumers.
Maintains a single instance of SimpleConsumer, periodically using the
consumer rebalancing algorithm to reassign partitions to this
SimpleConsumer.
"""
def __init__(self,
topic,
cluster,
consumer_group,
fetch_message_max_bytes=1024 * 1024,
num_consumer_fetchers=1,
auto_commit_enable=False,
auto_commit_interval_ms=60 * 1000,
queued_max_messages=2000,
fetch_min_bytes=1,
fetch_error_backoff_ms=500,
fetch_wait_max_ms=100,
offsets_channel_backoff_ms=1000,
offsets_commit_max_retries=5,
auto_offset_reset=OffsetType.EARLIEST,
consumer_timeout_ms=-1,
rebalance_max_retries=5,
rebalance_backoff_ms=2 * 1000,
zookeeper_connection_timeout_ms=6 * 1000,
zookeeper_connect=None,
zookeeper_hosts='127.0.0.1:2181',
zookeeper=None,
auto_start=True,
reset_offset_on_start=False,
post_rebalance_callback=None,
use_rdkafka=False,
compacted_topic=False,
membership_protocol=RangeProtocol,
deserializer=None,
reset_offset_on_fetch=True):
"""Create a BalancedConsumer instance
:param topic: The topic this consumer should consume
:type topic: :class:`pykafka.topic.Topic`
:param cluster: The cluster to which this consumer should connect
:type cluster: :class:`pykafka.cluster.Cluster`
:param consumer_group: The name of the consumer group this consumer
should join. Consumer group names are namespaced at the cluster level,
meaning that two consumers consuming different topics with the same group name
will be treated as part of the same group.
:type consumer_group: str
:param fetch_message_max_bytes: The number of bytes of messages to
attempt to fetch with each fetch request
:type fetch_message_max_bytes: int
:param num_consumer_fetchers: The number of workers used to make
FetchRequests
:type num_consumer_fetchers: int
:param auto_commit_enable: If true, periodically commit to kafka the
offset of messages already returned from consume() calls. Requires that
`consumer_group` is not `None`.
:type auto_commit_enable: bool
:param auto_commit_interval_ms: The frequency (in milliseconds) at which
the consumer's offsets are committed to kafka. This setting is
ignored if `auto_commit_enable` is `False`.
:type auto_commit_interval_ms: int
:param queued_max_messages: The maximum number of messages buffered for
consumption in the internal
:class:`pykafka.simpleconsumer.SimpleConsumer`
:type queued_max_messages: int
:param fetch_min_bytes: The minimum amount of data (in bytes) that the
server should return for a fetch request. If insufficient data is
available, the request will block until sufficient data is available.
:type fetch_min_bytes: int
:param fetch_error_backoff_ms: *UNUSED*.
See :class:`pykafka.simpleconsumer.SimpleConsumer`.
:type fetch_error_backoff_ms: int
:param fetch_wait_max_ms: The maximum amount of time (in milliseconds)
that the server will block before answering a fetch request if
there isn't sufficient data to immediately satisfy `fetch_min_bytes`.
:type fetch_wait_max_ms: int
:param offsets_channel_backoff_ms: Backoff time to retry failed offset
commits and fetches.
:type offsets_channel_backoff_ms: int
:param offsets_commit_max_retries: The number of times the offset commit
worker should retry before raising an error.
:type offsets_commit_max_retries: int
:param auto_offset_reset: What to do if an offset is out of range. This
setting indicates how to reset the consumer's internal offset
counter when an `OffsetOutOfRangeError` is encountered.
:type auto_offset_reset: :class:`pykafka.common.OffsetType`
:param consumer_timeout_ms: Amount of time (in milliseconds) the
consumer may spend without messages available for consumption
before returning None.
:type consumer_timeout_ms: int
:param rebalance_max_retries: The number of times the rebalance should
retry before raising an error.
:type rebalance_max_retries: int
:param rebalance_backoff_ms: Backoff time (in milliseconds) between
retries during rebalance.
:type rebalance_backoff_ms: int
:param zookeeper_connection_timeout_ms: The maximum time (in
milliseconds) that the consumer waits while establishing a
connection to zookeeper.
:type zookeeper_connection_timeout_ms: int
:param zookeeper_connect: Deprecated::2.7,3.6 Comma-Separated
(ip1:port1,ip2:port2) strings indicating the zookeeper nodes to which
to connect.
:type zookeeper_connect: str
:param zookeeper_hosts: KazooClient-formatted string of ZooKeeper hosts to which
to connect.
:type zookeeper_hosts: str
:param zookeeper: A KazooClient connected to a Zookeeper instance.
If provided, `zookeeper_connect` is ignored.
:type zookeeper: :class:`kazoo.client.KazooClient`
:param auto_start: Whether the consumer should begin communicating
with zookeeper after __init__ is complete. If false, communication
can be started with `start()`.
:type auto_start: bool
:param reset_offset_on_start: Whether the consumer should reset its
internal offset counter to `self._auto_offset_reset` and commit that
offset immediately upon starting up
:type reset_offset_on_start: bool
:param post_rebalance_callback: A function to be called when a rebalance is
in progress. This function should accept three arguments: the
:class:`pykafka.balancedconsumer.BalancedConsumer` instance that just
completed its rebalance, a dict of partitions that it owned before the
rebalance, and a dict of partitions it owns after the rebalance. These dicts
map partition ids to the most recently known offsets for those partitions.
This function can optionally return a dictionary mapping partition ids to
offsets. If it does, the consumer will reset its offsets to the supplied
values before continuing consumption.
Note that the BalancedConsumer is in a poorly defined state at
the time this callback runs, so that accessing its properties
(such as `held_offsets` or `partitions`) might yield confusing
results. Instead, the callback should really rely on the
provided partition-id dicts, which are well-defined.
:type post_rebalance_callback: function
:param use_rdkafka: Use librdkafka-backed consumer if available
:type use_rdkafka: bool
:param compacted_topic: Set to read from a compacted topic. Forces
consumer to use less stringent message ordering logic because compacted
topics do not provide offsets in strict incrementing order.
:type compacted_topic: bool
:param membership_protocol: The group membership protocol to which this consumer
should adhere
:type membership_protocol: :class:`pykafka.membershipprotocol.GroupMembershipProtocol`
:param deserializer: A function defining how to deserialize messages returned
from Kafka. A function with the signature d(value, partition_key) that
returns a tuple of (deserialized_value, deserialized_partition_key). The
arguments passed to this function are the bytes representations of a
message's value and partition key, and the returned data should be these
fields transformed according to the client code's serialization logic.
See `pykafka.utils.__init__` for stock implemtations.
:type deserializer: function
:param reset_offset_on_fetch: Whether to update offsets during fetch_offsets.
Disable for read-only use cases to prevent side-effects.
:type reset_offset_on_fetch: bool
"""
self._cluster = cluster
try:
self._consumer_group = get_string(consumer_group).encode('ascii')
except UnicodeEncodeError:
raise UnicodeException("Consumer group name '{}' contains non-ascii "
"characters".format(consumer_group))
self._topic = topic
self._auto_commit_enable = auto_commit_enable
self._auto_commit_interval_ms = valid_int(auto_commit_interval_ms)
self._fetch_message_max_bytes = valid_int(fetch_message_max_bytes)
self._fetch_min_bytes = valid_int(fetch_min_bytes)
self._rebalance_max_retries = valid_int(rebalance_max_retries, allow_zero=True)
self._num_consumer_fetchers = valid_int(num_consumer_fetchers)
self._queued_max_messages = valid_int(queued_max_messages)
self._fetch_wait_max_ms = valid_int(fetch_wait_max_ms, allow_zero=True)
self._rebalance_backoff_ms = valid_int(rebalance_backoff_ms)
self._consumer_timeout_ms = valid_int(consumer_timeout_ms,
allow_zero=True, allow_negative=True)
self._offsets_channel_backoff_ms = valid_int(offsets_channel_backoff_ms)
self._offsets_commit_max_retries = valid_int(offsets_commit_max_retries,
allow_zero=True)
self._auto_offset_reset = auto_offset_reset
self._zookeeper_connect = zookeeper_connect or zookeeper_hosts
self._zookeeper_connection_timeout_ms = valid_int(zookeeper_connection_timeout_ms,
allow_zero=True)
self._reset_offset_on_start = reset_offset_on_start
self._post_rebalance_callback = post_rebalance_callback
self._generation_id = -1
self._running = False
self._worker_exception = None
self._is_compacted_topic = compacted_topic
self._membership_protocol = membership_protocol
self._deserializer = deserializer
self._reset_offset_on_fetch = reset_offset_on_fetch
if not rdkafka and use_rdkafka:
raise ImportError("use_rdkafka requires rdkafka to be installed")
if GEventHandler and isinstance(self._cluster.handler, GEventHandler) and use_rdkafka:
raise ImportError("use_rdkafka cannot be used with gevent")
self._use_rdkafka = rdkafka and use_rdkafka
self._rebalancing_lock = cluster.handler.Lock()
self._rebalancing_in_progress = self._cluster.handler.Event()
self._internal_consumer_running = self._cluster.handler.Event()
self._consumer = None
self._consumer_id = get_bytes("{hostname}:{uuid}".format(
hostname=socket.gethostname(),
uuid=uuid4()
))
self._setting_watches = True
self._topic_path = '/consumers/{group}/owners/{topic}'.format(
group=get_string(self._consumer_group),
topic=self._topic.name)
self._consumer_id_path = '/consumers/{group}/ids'.format(
group=get_string(self._consumer_group))
self._zookeeper = None
self._owns_zookeeper = zookeeper is None
if zookeeper is not None:
self._zookeeper = zookeeper
if auto_start is True:
self.start()
def __del__(self):
log.debug("Finalising {}".format(self))
if self._running:
self.stop()
def __repr__(self):
return "<{module}.{name} at {id_} (consumer_group={group})>".format(
module=self.__class__.__module__,
name=self.__class__.__name__,
id_=hex(id(self)),
group=get_string(self._consumer_group)
)
def _raise_worker_exceptions(self):
"""Raises exceptions encountered on worker threads"""
if self._worker_exception is not None:
reraise(*self._worker_exception)
@property
def topic(self):
"""The topic this consumer consumes"""
return self._topic
@property
def partitions(self):
"""A list of the partitions that this consumer consumes"""
return self._consumer.partitions if self._consumer else dict()
@property
def _partitions(self):
"""Convenient shorthand for set of partitions internally held"""
return set(
[] if self.partitions is None else itervalues(self.partitions))
@property
def held_offsets(self):
"""Return a map from partition id to held offset for each partition"""
if not self._consumer:
return None
return self._consumer.held_offsets
def start(self):
"""Open connections and join a consumer group."""
try:
if self._zookeeper is None:
self._setup_zookeeper(self._zookeeper_connect,
self._zookeeper_connection_timeout_ms)
self._zookeeper.ensure_path(self._topic_path)
self._add_self()
self._running = True
self._set_watches()
self._rebalance()
except Exception:
log.exception("Stopping consumer in response to error")
self.stop()
def stop(self):
"""Close the zookeeper connection and stop consuming.
This method should be called as part of a graceful shutdown process.
"""
log.debug("Stopping {}".format(self))
with self._rebalancing_lock:
# We acquire the lock in order to prevent a race condition where a
# rebalance that is already underway might re-register the zk
# nodes that we remove here
self._running = False
if self._consumer is not None:
self._consumer.stop()
if self._owns_zookeeper:
# NB this should always come last, so we do not hand over control
# of our partitions until consumption has really been halted
self._zookeeper.stop()
else:
self._remove_partitions(self._get_held_partitions())
try:
self._zookeeper.delete(self._path_self)
except NoNodeException:
pass
# additionally we'd want to remove watches here, but there are no
# facilities for that in ChildrenWatch - as a workaround we check
# self._running in the watcher callbacks (see further down)
def _setup_zookeeper(self, zookeeper_connect, timeout):
"""Open a connection to a ZooKeeper host.
:param zookeeper_connect: The 'ip:port' address of the zookeeper node to
which to connect.
:type zookeeper_connect: str
:param timeout: Connection timeout (in milliseconds)
:type timeout: int
"""
kazoo_kwargs = {'timeout': timeout / 1000}
if GEventHandler and isinstance(self._cluster.handler, GEventHandler):
kazoo_kwargs['handler'] = SequentialGeventHandler()
self._zookeeper = KazooClient(zookeeper_connect, **kazoo_kwargs)
self._zookeeper.start()
def _setup_internal_consumer(self, partitions=None, start=True):
"""Instantiate an internal SimpleConsumer instance"""
if partitions is None:
partitions = []
# Only re-create internal consumer if something changed.
if partitions != self._partitions:
cns = self._get_internal_consumer(partitions=list(partitions), start=start)
if self._post_rebalance_callback is not None:
old_offsets = (self._consumer.held_offsets
if self._consumer else dict())
new_offsets = cns.held_offsets
try:
reset_offsets = self._post_rebalance_callback(
self, old_offsets, new_offsets)
except Exception:
log.exception("post rebalance callback threw an exception")
self._worker_exception = sys.exc_info()
return False
if reset_offsets:
cns.reset_offsets(partition_offsets=[
(cns.partitions[id_], offset) for
(id_, offset) in iteritems(reset_offsets)])
self._consumer = cns
if self._consumer and self._consumer._running:
if not self._internal_consumer_running.is_set():
self._internal_consumer_running.set()
else:
if self._internal_consumer_running.is_set():
self._internal_consumer_running.clear()
return True
def _get_internal_consumer(self, partitions=None, start=True):
"""Instantiate a SimpleConsumer for internal use.
If there is already a SimpleConsumer instance held by this object,
disable its workers and mark it for garbage collection before
creating a new one.
"""
if partitions is None:
partitions = []
reset_offset_on_start = self._reset_offset_on_start
if self._consumer is not None:
self._consumer.stop()
# only use this setting for the first call to
# _get_internal_consumer. subsequent calls should not
# reset the offsets, since they can happen at any time
reset_offset_on_start = False
Cls = (rdkafka.RdKafkaSimpleConsumer
if self._use_rdkafka else SimpleConsumer)
cns = Cls(
self._topic,
self._cluster,
consumer_group=self._consumer_group,
partitions=partitions,
auto_commit_enable=self._auto_commit_enable,
auto_commit_interval_ms=self._auto_commit_interval_ms,
fetch_message_max_bytes=self._fetch_message_max_bytes,
fetch_min_bytes=self._fetch_min_bytes,
num_consumer_fetchers=self._num_consumer_fetchers,
queued_max_messages=self._queued_max_messages,
fetch_wait_max_ms=self._fetch_wait_max_ms,
consumer_timeout_ms=self._consumer_timeout_ms,
offsets_channel_backoff_ms=self._offsets_channel_backoff_ms,
offsets_commit_max_retries=self._offsets_commit_max_retries,
auto_offset_reset=self._auto_offset_reset,
reset_offset_on_start=reset_offset_on_start,
auto_start=False,
compacted_topic=self._is_compacted_topic,
deserializer=self._deserializer,
reset_offset_on_fetch=self._reset_offset_on_fetch
)
cns.consumer_id = self._consumer_id
cns.generation_id = self._generation_id
if start:
cns.start()
return cns
def _get_participants(self):
"""Use zookeeper to get the other consumers of this topic.
:return: A sorted list of the ids of other consumers of this
consumer's topic
"""
try:
consumer_ids = self._zookeeper.get_children(self._consumer_id_path)
except NoNodeException:
log.debug("Consumer group doesn't exist. "
"No participants to find")
return []
participants = []
for id_ in consumer_ids:
try:
topic, stat = self._zookeeper.get("%s/%s" % (self._consumer_id_path, id_))
if topic == self._topic.name:
participants.append(get_bytes(id_))
except NoNodeException:
pass # node disappeared between ``get_children`` and ``get``
participants = sorted(participants)
return participants
def _build_watch_callback(self, fn, proxy):
"""Return a function that's safe to use as a ChildrenWatch callback
Fixes the issue from https://github.com/Parsely/pykafka/issues/345
"""
def _callback(children):
# discover whether the referenced object still exists
try:
proxy.__repr__()
except ReferenceError:
return False
return fn(proxy, children)
return _callback
def _set_watches(self):
"""Set watches in zookeeper that will trigger rebalances.
Rebalances should be triggered whenever a broker, topic, or consumer
znode is changed in zookeeper. This ensures that the balance of the
consumer group remains up-to-date with the current state of the
cluster.
"""
proxy = weakref.proxy(self)
_brokers_changed = self._build_watch_callback(BalancedConsumer._brokers_changed, proxy)
_topics_changed = self._build_watch_callback(BalancedConsumer._topics_changed, proxy)
_consumers_changed = self._build_watch_callback(BalancedConsumer._consumers_changed, proxy)
self._setting_watches = True
# Set all our watches and then rebalance
broker_path = '/brokers/ids'
try:
self._broker_watcher = ChildrenWatch(
self._zookeeper, broker_path,
_brokers_changed
)
except NoNodeException:
raise Exception(
'The broker_path "%s" does not exist in your '
'ZooKeeper cluster -- is your Kafka cluster running?'
% broker_path)
self._topics_watcher = ChildrenWatch(
self._zookeeper,
'/brokers/topics',
_topics_changed
)
self._consumer_watcher = ChildrenWatch(
self._zookeeper, self._consumer_id_path,
_consumers_changed
)
self._setting_watches = False
def _add_self(self):
"""Register this consumer in zookeeper."""
self._zookeeper.create(
self._path_self, self._topic.name, ephemeral=True, makepath=True)
@property
def _path_self(self):
"""Path where this consumer should be registered in zookeeper"""
return '{path}/{id_}'.format(
path=self._consumer_id_path,
# get_string is necessary to avoid writing literal "b'" to zookeeper
id_=get_string(self._consumer_id)
)
def _update_member_assignment(self):
"""Decide and assign new partitions for this consumer"""
for i in range(self._rebalance_max_retries):
try:
# If retrying, be sure to make sure the
# partition allocation is correct.
participants = self._get_participants()
if self._consumer_id not in participants:
# situation that only occurs if our zk session expired
self._add_self()
participants.append(self._consumer_id)
new_partitions = self._membership_protocol.decide_partitions(
participants, self._topic.partitions, self._consumer_id)
if not new_partitions:
log.warning("No partitions assigned to consumer %s",
self._consumer_id)
# Update zk with any changes:
# Note that we explicitly fetch our set of held partitions
# from zk, rather than assuming it will be identical to
# `self.partitions`. This covers the (rare) situation
# where due to an interrupted connection our zk session
# has expired, in which case we'd hold zero partitions on
# zk, but `self._partitions` may be outdated and non-empty
current_zk_parts = self._get_held_partitions()
self._remove_partitions(current_zk_parts - new_partitions)
self._add_partitions(new_partitions - current_zk_parts)
if self._setup_internal_consumer(new_partitions):
log.info('Rebalancing Complete.')
break
except PartitionOwnedError as ex:
if i == self._rebalance_max_retries - 1:
log.warning('Failed to acquire partition %s after %d retries.',
ex.partition, i)
raise
log.info('Unable to acquire partition %s. Retrying', ex.partition)
self._cluster.handler.sleep(i * (self._rebalance_backoff_ms / 1000))
def _rebalance(self):
"""Start the rebalancing process for this consumer
This method is called whenever a zookeeper watch is triggered.
"""
# This Event is used to notify about rebalance operation to SimpleConsumer's consume().
if not self._rebalancing_in_progress.is_set():
self._rebalancing_in_progress.set()
if self._consumer is not None:
self.commit_offsets()
# this is necessary because we can't stop() while the lock is held
# (it's not an RLock)
with self._rebalancing_lock:
if not self._running:
raise ConsumerStoppedException
log.info('Rebalancing consumer "%s" for topic "%s".' % (
self._consumer_id, self._topic.name))
self._update_member_assignment()
if self._rebalancing_in_progress.is_set():
self._rebalancing_in_progress.clear()
def _path_from_partition(self, p):
"""Given a partition, return its path in zookeeper.
:type p: :class:`pykafka.partition.Partition`
"""
return "%s/%s-%s" % (self._topic_path, p.leader.id, p.id)
def _remove_partitions(self, partitions):
"""Remove partitions from the zookeeper registry for this consumer.
:param partitions: The partitions to remove.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
try:
# TODO pass zk node version to make sure we still own this node
self._zookeeper.delete(self._path_from_partition(p))
except NoNodeException:
pass
def _add_partitions(self, partitions):
"""Add partitions to the zookeeper registry for this consumer.
:param partitions: The partitions to add.
:type partitions: Iterable of :class:`pykafka.partition.Partition`
"""
for p in partitions:
try:
self._zookeeper.create(
self._path_from_partition(p),
value=get_bytes(self._consumer_id),
ephemeral=True
)
except NodeExistsError:
raise PartitionOwnedError(p)
def _get_held_partitions(self):
"""Build a set of partitions zookeeper says we own"""
zk_partition_ids = set()
all_partitions = self._zookeeper.get_children(self._topic_path)
for partition_slug in all_partitions:
try:
owner_id, stat = self._zookeeper.get(
'{path}/{slug}'.format(
path=self._topic_path, slug=partition_slug))
if owner_id == get_bytes(self._consumer_id):
zk_partition_ids.add(int(partition_slug.split('-')[1]))
except NoNodeException:
pass # disappeared between ``get_children`` and ``get``
return set(self._topic.partitions[_id] for _id in zk_partition_ids)
@_catch_thread_exception
def _brokers_changed(self, brokers):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by broker change ({})".format(
self._consumer_id))
self._rebalance()
@_catch_thread_exception
def _consumers_changed(self, consumers):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by consumer change ({})".format(
self._consumer_id))
self._rebalance()
@_catch_thread_exception
def _topics_changed(self, topics):
if not self._running:
return False # `False` tells ChildrenWatch to disable this watch
if self._setting_watches:
return
log.debug("Rebalance triggered by topic change ({})".format(
self._consumer_id))
self._rebalance()
def reset_offsets(self, partition_offsets=None):
"""Reset offsets for the specified partitions
For each value provided in `partition_offsets`: if the value is an integer,
immediately reset the partition's internal offset counter to that value. If
it's a `datetime.datetime` instance or a valid `OffsetType`, issue a
`ListOffsetRequest` using that timestamp value to discover the latest offset
in the latest log segment before that timestamp, then set the partition's
internal counter to that value.
:param partition_offsets: (`partition`, `timestamp_or_offset`) pairs to
reset where `partition` is the partition for which to reset the offset
and `timestamp_or_offset` is EITHER the timestamp before which to find
a valid offset to set the partition's counter to OR the new offset the
partition's counter should be set to
:type partition_offsets: Sequence of tuples of the form
(:class:`pykafka.partition.Partition`, int OR `datetime.datetime`)
"""
self._raise_worker_exceptions()
if not self._consumer:
raise ConsumerStoppedException("Internal consumer is stopped")
self._consumer.reset_offsets(partition_offsets=partition_offsets)
def consume(self, block=True):
"""Get one message from the consumer
:param block: Whether to block while waiting for a message
:type block: bool
"""
def consumer_timed_out():
"""Indicates whether the consumer has received messages recently"""
if self._consumer_timeout_ms == -1:
return False
disp = (time.time() - self._last_message_time) * 1000.0
return disp > self._consumer_timeout_ms
message = None
self._last_message_time = time.time()
while message is None and not consumer_timed_out():
if not self._internal_consumer_running.is_set():
self._cluster.handler.sleep()
self._raise_worker_exceptions()
self._internal_consumer_running.wait(self._consumer_timeout_ms / 1000)
try:
# acquire the lock to ensure that we don't start trying to consume from
# a _consumer that might soon be replaced by an in-progress rebalance
with self._rebalancing_lock:
message = self._consumer.consume(block=block, unblock_event=self._rebalancing_in_progress)
# If Gevent is used, waiting to acquire _rebalancing lock introduces a race condition.
# This sleep would ensure that the _rebalance method acquires the _rebalancing_lock
# Issue: https://github.com/Parsely/pykafka/issues/671
if self._rebalancing_in_progress.is_set():
self._cluster.handler.sleep()
except (ConsumerStoppedException, AttributeError):
if not self._running:
raise ConsumerStoppedException
if message:
self._last_message_time = time.time()
if not block:
return message
return message
def __iter__(self):
"""Yield an infinite stream of messages until the consumer times out"""
while True:
message = self.consume(block=True)
if not message:
return
yield message
def commit_offsets(self, partition_offsets=None):
"""Commit offsets for this consumer's partitions
Uses the offset commit/fetch API
:param partition_offsets: (`partition`, `offset`) pairs to
commit where `partition` is the partition for which to commit the offset
and `offset` is the offset to commit for the partition. Note that using
this argument when `auto_commit_enable` is enabled can cause inconsistencies
in committed offsets. For best results, use *either* this argument *or*
`auto_commit_enable`.
:type partition_offsets: Sequence of tuples of the form
(:class:`pykafka.partition.Partition`, int)
"""
self._raise_worker_exceptions()
if not self._consumer:
raise KafkaException("Cannot commit offsets - consumer not started")
return self._consumer.commit_offsets(partition_offsets=partition_offsets)
class ZooKeeper(AbstractDCS):
def __init__(self, config):
super(ZooKeeper, self).__init__(config)
hosts = config.get('hosts', [])
if isinstance(hosts, list):
hosts = ','.join(hosts)
self._client = KazooClient(hosts,
handler=PatroniSequentialThreadingHandler(
config['retry_timeout']),
timeout=config['ttl'],
connection_retry={
'max_delay': 1,
'max_tries': -1
},
command_retry={
'deadline': config['retry_timeout'],
'max_delay': 1,
'max_tries': -1
})
self._client.add_listener(self.session_listener)
self._my_member_data = None
self._fetch_cluster = True
self._orig_kazoo_connect = self._client._connection._connect
self._client._connection._connect = self._kazoo_connect
self._client.start()
def _kazoo_connect(self, host, port):
"""Kazoo is using Ping's to determine health of connection to zookeeper. If there is no
response on Ping after Ping interval (1/2 from read_timeout) it will consider current
connection dead and try to connect to another node. Without this "magic" it was taking
up to 2/3 from session timeout (ttl) to figure out that connection was dead and we had
only small time for reconnect and retry.
This method is needed to return different value of read_timeout, which is not calculated
from negotiated session timeout but from value of `loop_wait`. And it is 2 sec smaller
than loop_wait, because we can spend up to 2 seconds when calling `touch_member()` and
`write_leader_optime()` methods, which also may hang..."""
ret = self._orig_kazoo_connect(host, port)
return max(self.loop_wait - 2, 2) * 1000, ret[1]
def session_listener(self, state):
if state in [KazooState.SUSPENDED, KazooState.LOST]:
self.cluster_watcher(None)
def cluster_watcher(self, event):
self._fetch_cluster = True
self.event.set()
def reload_config(self, config):
self.set_retry_timeout(config['retry_timeout'])
loop_wait = config['loop_wait']
loop_wait_changed = self._loop_wait != loop_wait
self._loop_wait = loop_wait
self._client.handler.set_connect_timeout(loop_wait)
# We need to reestablish connection to zookeeper if we want to change
# read_timeout (and Ping interval respectively), because read_timeout
# is calculated in `_kazoo_connect` method. If we are changing ttl at
# the same time, set_ttl method will reestablish connection and return
# `!True`, otherwise we will close existing connection and let kazoo
# open the new one.
if not self.set_ttl(int(config['ttl'] * 1000)) and loop_wait_changed:
self._client._connection._socket.close()
def set_ttl(self, ttl):
"""It is not possible to change ttl (session_timeout) in zookeeper without
destroying old session and creating the new one. This method returns `!True`
if session_timeout has been changed (`restart()` has been called)."""
if self._client._session_timeout != ttl:
self._client._session_timeout = ttl
self._client.restart()
return True
def set_retry_timeout(self, retry_timeout):
self._client._retry.deadline = retry_timeout
def get_node(self, key, watch=None):
try:
ret = self._client.get(key, watch)
return (ret[0].decode('utf-8'), ret[1])
except NoNodeError:
return None
@staticmethod
def member(name, value, znode):
return Member.from_node(znode.version, name, znode.ephemeralOwner,
value)
def get_children(self, key, watch=None):
try:
return self._client.get_children(key, watch)
except NoNodeError:
return []
def load_members(self):
members = []
for member in self.get_children(self.members_path,
self.cluster_watcher):
data = self.get_node(self.members_path + member)
if data is not None:
members.append(self.member(member, *data))
return members
def _inner_load_cluster(self):
self._fetch_cluster = False
self.event.clear()
nodes = set(
self.get_children(self.client_path(''), self.cluster_watcher))
if not nodes:
self._fetch_cluster = True
# get initialize flag
initialize = (self.get_node(self.initialize_path)
or [None])[0] if self._INITIALIZE in nodes else None
# get global dynamic configuration
config = self.get_node(
self.config_path,
watch=self.cluster_watcher) if self._CONFIG in nodes else None
config = config and ClusterConfig.from_node(config[1].version,
config[0], config[1].mzxid)
# get last leader operation
last_leader_operation = self._OPTIME in nodes and self._fetch_cluster and self.get_node(
self.leader_optime_path)
last_leader_operation = last_leader_operation and int(
last_leader_operation[0]) or 0
# get list of members
members = self.load_members() if self._MEMBERS[:-1] in nodes else []
# get leader
leader = self.get_node(
self.leader_path) if self._LEADER in nodes else None
if leader:
client_id = self._client.client_id
if not self._ctl and leader[0] == self._name and client_id is not None \
and client_id[0] != leader[1].ephemeralOwner:
logger.info(
'I am leader but not owner of the session. Removing leader node'
)
self._client.delete(self.leader_path)
leader = None
if leader:
member = Member(-1, leader[0], None, {})
member = ([m for m in members if m.name == leader[0]]
or [member])[0]
leader = Leader(leader[1].version, leader[1].ephemeralOwner,
member)
self._fetch_cluster = member.index == -1
# failover key
failover = self.get_node(
self.failover_path,
watch=self.cluster_watcher) if self._FAILOVER in nodes else None
failover = failover and Failover.from_node(failover[1].version,
failover[0])
self._cluster = Cluster(initialize, config, leader,
last_leader_operation, members, failover)
def _load_cluster(self):
if self._fetch_cluster or self._cluster is None:
try:
self._client.retry(self._inner_load_cluster)
except Exception:
logger.exception('get_cluster')
self.cluster_watcher(None)
raise ZooKeeperError('ZooKeeper in not responding properly')
def _create(self, path, value, **kwargs):
try:
self._client.retry(self._client.create, path,
value.encode('utf-8'), **kwargs)
return True
except:
return False
def attempt_to_acquire_leader(self, permanent=False):
ret = self._create(self.leader_path,
self._name,
makepath=True,
ephemeral=not permanent)
if not ret:
logger.info('Could not take out TTL lock')
return ret
def set_failover_value(self, value, index=None):
try:
self._client.retry(self._client.set,
self.failover_path,
value.encode('utf-8'),
version=index or -1)
return True
except NoNodeError:
return value == '' or (index is None
and self._create(self.failover_path, value))
except:
logging.exception('set_failover_value')
return False
def set_config_value(self, value, index=None):
try:
self._client.retry(self._client.set,
self.config_path,
value.encode('utf-8'),
version=index or -1)
return True
except NoNodeError:
return index is None and self._create(self.config_path, value)
except Exception:
logging.exception('set_config_value')
return False
def initialize(self, create_new=True, sysid=""):
return self._create(self.initialize_path, sysid, makepath=True) if create_new \
else self._client.retry(self._client.set, self.initialize_path, sysid.encode("utf-8"))
def touch_member(self, data, ttl=None, permanent=False):
cluster = self.cluster
member = cluster and cluster.get_member(self._name,
fallback_to_leader=False)
data = data.encode('utf-8')
if member and self._client.client_id is not None and member.session != self._client.client_id[
0]:
try:
self._client.delete_async(self.member_path).get(timeout=1)
except NoNodeError:
pass
except:
return False
member = None
if member:
if data == self._my_member_data:
return True
else:
try:
self._client.create_async(
self.member_path,
data,
makepath=True,
ephemeral=not permanent).get(timeout=1)
self._my_member_data = data
return True
except Exception as e:
if not isinstance(e, NodeExistsError):
logger.exception('touch_member')
return False
try:
self._client.set_async(self.member_path, data).get(timeout=1)
self._my_member_data = data
return True
except:
logger.exception('touch_member')
return False
def take_leader(self):
return self.attempt_to_acquire_leader()
def _write_leader_optime(self, last_operation):
last_operation = last_operation.encode('utf-8')
try:
self._client.set_async(self.leader_optime_path,
last_operation).get(timeout=1)
return True
except NoNodeError:
try:
self._client.create_async(self.leader_optime_path,
last_operation,
makepath=True).get(timeout=1)
return True
except:
logger.exception('Failed to create %s',
self.leader_optime_path)
except:
logger.exception('Failed to update %s', self.leader_optime_path)
return False
def update_leader(self):
return True
def delete_leader(self):
self._client.restart()
self._my_member_data = None
return True
def _cancel_initialization(self):
node = self.get_node(self.initialize_path)
if node:
self._client.delete(self.initialize_path, version=node[1].version)
def cancel_initialization(self):
try:
self._client.retry(self._cancel_initialization)
except:
logger.exception("Unable to delete initialize key")
def delete_cluster(self):
try:
return self._client.retry(self._client.delete,
self.client_path(''),
recursive=True)
except NoNodeError:
return True
def watch(self, timeout):
if super(ZooKeeper, self).watch(timeout):
self._fetch_cluster = True
return self._fetch_cluster
class ZkStateManager(StateManager):
"""
State manager which connects to zookeeper and
gets and sets states from there.
"""
def __init__(self, name, host, port, rootpath, tunnelhost):
self.name = name
self.host = host
self.port = port
self.tunnelhost = tunnelhost
self.rootpath = rootpath
def start(self):
""" state Zookeeper """
if self.is_host_port_reachable():
self.client = KazooClient(self.hostport)
else:
localport = self.establish_ssh_tunnel()
self.client = KazooClient("localhost:" + str(localport))
self.client.start()
def on_connection_change(state):
""" callback to log """
LOG.info("Connection state changed to: " + state)
self.client.add_listener(on_connection_change)
def stop(self):
""" stop Zookeeper """
self.client.stop()
self.terminate_ssh_tunnel()
# pylint: disable=function-redefined
def get_topologies(self, callback=None):
""" get topologies """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""Custom callback to get the topologies right now."""
ret["result"] = data
try:
self._get_topologies_with_watch(callback, isWatching)
except NoNodeError as err:
self.client.stop()
path = self.get_topologies_path()
raise StateException("Error required topology path '%s' not found" % (path),
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
# The topologies are now populated with the data.
return ret["result"]
def _get_topologies_with_watch(self, callback, isWatching):
"""
Helper function to get topologies with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_topologies_path()
if isWatching:
LOG.info("Adding children watch for path: " + path)
# pylint: disable=unused-variable
@self.client.ChildrenWatch(path)
def watch_topologies(topologies):
""" callback to watch topologies """
callback(topologies)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def get_topology(self, topologyName, callback=None):
""" get topologies """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""Custom callback to get the topologies right now."""
ret["result"] = data
self._get_topology_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_topology_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_topology_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_topology(data, stats):
""" watch topology """
if data:
topology = Topology()
topology.ParseFromString(data)
callback(topology)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_topology(self, topologyName, topology):
""" crate topology """
if not topology or not topology.IsInitialized():
raise StateException("Topology protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_topology_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
topologyString = topology.SerializeToString()
try:
self.client.create(path, value=topologyString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating topology",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating topology",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating topology",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_topology(self, topologyName):
""" delete topology """
path = self.get_topology_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deteling topology",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting topology",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting topology",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_pplan(self, topologyName, callback=None):
""" get physical plan """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_pplan_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_pplan_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_pplan_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_pplan(data, stats):
""" invoke callback to watch physical plan """
if data:
pplan = PhysicalPlan()
pplan.ParseFromString(data)
callback(pplan)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_pplan(self, topologyName, pplan):
""" create physical plan """
if not pplan or not pplan.IsInitialized():
raise StateException("Physical Plan protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_pplan_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
pplanString = pplan.SerializeToString()
try:
self.client.create(path, value=pplanString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating pplan",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating pplan",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating pplan",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_pplan(self, topologyName):
""" delete physical plan info """
path = self.get_pplan_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deleting pplan",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting pplan",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting pplan",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_execution_state(self, topologyName, callback=None):
""" get execution state """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_execution_state_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_execution_state_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get execution state with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_execution_state_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_execution_state(data, stats):
""" invoke callback to watch execute state """
if data:
executionState = ExecutionState()
executionState.ParseFromString(data)
callback(executionState)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def create_execution_state(self, topologyName, executionState):
""" create execution state """
if not executionState or not executionState.IsInitialized():
raise StateException("Execution State protobuf not init properly",
StateException.EX_TYPE_PROTOBUF_ERROR), None, sys.exc_info()[2]
path = self.get_execution_state_path(topologyName)
LOG.info("Adding topology: {0} to path: {1}".format(
topologyName, path))
executionStateString = executionState.SerializeToString()
try:
self.client.create(path, value=executionStateString, makepath=True)
return True
except NoNodeError:
raise StateException("NoNodeError while creating execution state",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NodeExistsError:
raise StateException("NodeExistsError while creating execution state",
StateException.EX_TYPE_NODE_EXISTS_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while creating execution state",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def delete_execution_state(self, topologyName):
""" delete execution state """
path = self.get_execution_state_path(topologyName)
LOG.info("Removing topology: {0} from path: {1}".format(
topologyName, path))
try:
self.client.delete(path)
return True
except NoNodeError:
raise StateException("NoNodeError while deleting execution state",
StateException.EX_TYPE_NO_NODE_ERROR), None, sys.exc_info()[2]
except NotEmptyError:
raise StateException("NotEmptyError while deleting execution state",
StateException.EX_TYPE_NOT_EMPTY_ERROR), None, sys.exc_info()[2]
except ZookeeperError:
raise StateException("Zookeeper while deleting execution state",
StateException.EX_TYPE_ZOOKEEPER_ERROR), None, sys.exc_info()[2]
except Exception:
# Just re raise the exception.
raise
def get_tmaster(self, topologyName, callback=None):
""" get tmaster """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the topologies right now.
"""
ret["result"] = data
self._get_tmaster_with_watch(topologyName, callback, isWatching)
# The topologies are now populated with the data.
return ret["result"]
def _get_tmaster_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get pplan with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_tmaster_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_tmaster(data, stats):
""" invoke callback to watch tmaster """
if data:
tmaster = TMasterLocation()
tmaster.ParseFromString(data)
callback(tmaster)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
def get_scheduler_location(self, topologyName, callback=None):
""" get scheduler location """
isWatching = False
# Temp dict used to return result
# if callback is not provided.
ret = {
"result": None
}
if callback:
isWatching = True
else:
def callback(data):
"""
Custom callback to get the scheduler location right now.
"""
ret["result"] = data
self._get_scheduler_location_with_watch(topologyName, callback, isWatching)
return ret["result"]
def _get_scheduler_location_with_watch(self, topologyName, callback, isWatching):
"""
Helper function to get scheduler location with
a callback. The future watch is placed
only if isWatching is True.
"""
path = self.get_scheduler_location_path(topologyName)
if isWatching:
LOG.info("Adding data watch for path: " + path)
# pylint: disable=unused-variable, unused-argument
@self.client.DataWatch(path)
def watch_scheduler_location(data, stats):
""" invoke callback to watch scheduler location """
if data:
scheduler_location = SchedulerLocation()
scheduler_location.ParseFromString(data)
callback(scheduler_location)
else:
callback(None)
# Returning False will result in no future watches
# being triggered. If isWatching is True, then
# the future watches will be triggered.
return isWatching
class Zk(object):
def __init__(self, run_id, conf):
self.run_id = run_id
self.conf = conf
self._zk = KazooClient(hosts=metadata.zk)
self._zk.start()
def clean(self):
#if self._zk.exists(metadata.topics_path):
# self._zk.delete(metadata.topics_path,recursive=True)
#if self._zk.exists(metadata.leader_path):
# self._zk.delete(metadata.leader_path,recursive=True)
#if self._zk.exists(metadata.rb_path):
# self._zk.delete(metadata.rb_path,recursive=True)
if self._zk.exists(metadata.experiment_path):
self._zk.delete(metadata.experiment_path, recursive=True)
def stop(self):
self._zk.stop()
def setup(self):
self.create_paths()
self.install_watches()
def wait(self, barrier):
if (barrier == 'subscriber'):
self.sub_barrier.wait()
elif (barrier == 'finished'):
self.finished_barrier.wait()
elif (barrier == 'monitoring'):
self.monitoring_barrier.wait()
else:
print('invalid barrier name')
def create_paths(self):
#create zk path for subscribers
for client in self.conf.client_numSubscribers.keys():
sub_path='%s/%s/sub/region_%s/%s'%\
(metadata.experiment_path,self.run_id,client[3:client.index('-')],client)
self._zk.ensure_path(sub_path)
#create zk path for publishers
for client in self.conf.client_numPublishers.keys():
pub_path='%s/%s/pub/region_%s/%s'%\
(metadata.experiment_path,self.run_id,client[3:client.index('-')],client)
self._zk.ensure_path(pub_path)
#create zk path to track joined subscribers and publishers
joined_sub_path='%s/%s/joined_sub'%\
(metadata.experiment_path,self.run_id)
self._zk.ensure_path(joined_sub_path)
joined_pub_path='%s/%s/joined_pub'%\
(metadata.experiment_path,self.run_id)
self._zk.ensure_path(joined_pub_path)
#create zk path to track subscribers have left
#self._zk.ensure_path('%s/%s/left_sub'%(metadata.experiment_path,self.run_id))
#create zk path to track monitoring processes have joined
eb_monitoring_path = '%s/%s/monitoring/eb' % (metadata.experiment_path,
self.run_id)
self._zk.ensure_path(eb_monitoring_path)
self.eb_monitors_exited = False
rb_monitoring_path = '%s/%s/monitoring/rb' % (metadata.experiment_path,
self.run_id)
self._zk.ensure_path(rb_monitoring_path)
self.rb_monitors_exited = False
#create barrier paths
sub_barrier_path = '%s/%s/barriers/sub' % (metadata.experiment_path,
self.run_id)
self._zk.ensure_path(sub_barrier_path)
pub_barrier_path = '%s/%s/barriers/pub' % (metadata.experiment_path,
self.run_id)
self._zk.ensure_path(pub_barrier_path)
finished_barrier_path = '%s/%s/barriers/finished' % (
metadata.experiment_path, self.run_id)
self._zk.ensure_path(finished_barrier_path)
monitoring_barrier_path = '%s/%s/barriers/monitoring' % (
metadata.experiment_path, self.run_id)
self._zk.ensure_path(monitoring_barrier_path)
#create barriers
self.sub_barrier = Barrier(client=self._zk, path=sub_barrier_path)
self.finished_barrier = Barrier(client=self._zk,
path=finished_barrier_path)
self.pub_barrier = Barrier(client=self._zk, path=pub_barrier_path)
self.monitoring_barrier = Barrier(client=self._zk,
path=monitoring_barrier_path)
def install_watches(self):
region_joined_subscriber_clients={region:0 for region in \
self.conf.region_clientsubscribers_map.keys()}
region_joined_publisher_clients={region:0 for region in \
self.conf.region_clientpublishers_map.keys()}
def _joined_endpoint_listener(children, event):
if event and event.type == EventType.CHILD:
if 'sub' in event.path:
client = event.path.rpartition('/')[2]
region = client[3:client.index('-')]
if (len(children) ==
self.conf.client_numSubscribers[client]):
print('All subscribers have joined on client:%s\n' %
(client))
region_joined_subscriber_clients[region] += 1
if (region_joined_subscriber_clients[region]==\
self.conf.region_clientsubscribers_map[region]):
self._zk.ensure_path('%s/%s/joined_sub/region_%s'%\
(metadata.experiment_path,self.run_id,region))
if (len(children) == 0):
print('All subscribers on client:%s have exited\n' %
(client))
region_joined_subscriber_clients[region] -= 1
if (region_joined_subscriber_clients[region] == 0):
self._zk.delete('%s/%s/joined_sub/region_%s'\
%(metadata.experiment_path,self.run_id,region))
return False
if 'pub' in event.path:
client = event.path.rpartition('/')[2]
region = client[3:client.index('-')]
if (len(children) == self.conf.client_numPublishers[client]
):
print('All publishers have joined on client:%s\n' %
(client))
region_joined_publisher_clients[region] += 1
if (region_joined_publisher_clients[region]==\
self.conf.region_clientpublishers_map[region]):
self._zk.ensure_path('%s/%s/joined_pub/region_%s'%\
(metadata.experiment_path,self.run_id,region))
return False
def _open_barrier(children, event):
if event and event.type == EventType.CHILD:
if 'joined_sub' in event.path:
if (len(children) == len(
self.conf.region_clientsubscribers_map)):
print(
"All subscribers have joined. Opening subscriber barrier\n"
)
self.sub_barrier.remove()
if (len(children) == 0):
print(
"All subscribers have left. Opening finished barrier\n"
)
self.finished_barrier.remove()
return False
if 'joined_pub' in event.path:
if (len(children) == len(
self.conf.region_clientpublishers_map)):
print(
"All publishers have joined. Opening publisher barrier\n"
)
self.pub_barrier.remove()
return False
if 'monitoring/eb' in event.path:
if (len(children) == 0):
print('All eb monitors have exited')
self.eb_monitors_exited = True
if (self.eb_monitors_exited
and self.rb_monitors_exited):
print(
'All monitors have exited. Opening monitoring barrier'
)
self.monitoring_barrier.remove()
return False
if 'monitoring/rb' in event.path:
if (len(children) == 0):
print('All rb monitors have exited')
self.rb_monitors_exited = True
if (self.eb_monitors_exited
and self.rb_monitors_exited):
print(
'All monitors have exited. Opening monitoring barrier'
)
self.monitoring_barrier.remove()
return False
sub_barrier_opener_watch=ChildrenWatch(client=self._zk,\
path='%s/%s/joined_sub'%(metadata.experiment_path,self.run_id),\
func=_open_barrier,send_event=True)
pub_barrier_opener_watch=ChildrenWatch(client=self._zk,\
path='%s/%s/joined_pub'%(metadata.experiment_path,self.run_id),\
func=_open_barrier,send_event=True)
eb_monitoring_watch=ChildrenWatch(client=self._zk,\
path='%s/%s/monitoring/eb'%(metadata.experiment_path,self.run_id),\
func=_open_barrier,send_event=True)
rb_monitoring_watch=ChildrenWatch(client=self._zk,\
path='%s/%s/monitoring/rb'%(metadata.experiment_path,self.run_id),\
func=_open_barrier,send_event=True)
joined_sub_watches=[ChildrenWatch(client=self._zk,\
path='%s/%s/sub/region_%s/%s'%\
(metadata.experiment_path,self.run_id,client[3:client.index('-')],client),\
func=_joined_endpoint_listener,send_event=True)\
for client in self.conf.client_numSubscribers.keys()]
joined_pub_watches=[ChildrenWatch(client=self._zk,\
path='%s/%s/pub/region_%s/%s'%\
(metadata.experiment_path,self.run_id,client[3:client.index('-')],client),\
func=_joined_endpoint_listener,send_event=True)\
for client in self.conf.client_numPublishers.keys()]
class TestServiceDiscovery(unittest.TestCase):
def setUp(self):
self.maxDiff = None
logging.basicConfig(format="%(asctime)s %(levelname)s %(module)s[%(lineno)d] %(threadName)s %(message)s", level=logging.WARN)
self.log = logging.getLogger()
self.basePath = "/discovery_test_%x" % int(time.time())
self.log.info("Using base path: %s" % self.basePath)
self.client = KazooClient(hosts="127.0.0.1:2181")
self.client.start()
self._clean()
self.discovery = ServiceDiscovery(self.client, self.basePath)
def _clean(self):
try:
self.client.delete(self.basePath,recursive=True)
except NoNodeError:
pass
def tearDown(self):
self.discovery.close()
self._clean()
self.client.stop()
self.client.close()
def test_paths(self):
svc1 = ServiceInstance.builder().id("instance1").name("service1").build()
svc2 = ServiceInstance.builder().id("instance2").name("service1").build()
svc3 = ServiceInstance.builder().id("foo1").name("foo").build()
self.assertEquals(self.basePath + "/service1/instance1", self.discovery.pathForInstance(svc1.getName(),svc1.getId()))
self.assertEquals(self.basePath + "/service1/instance2", self.discovery.pathForInstance(svc2.getName(),svc2.getId()))
self.assertEquals(self.basePath + "/foo/foo1", self.discovery.pathForInstance(svc3.getName(),svc3.getId()))
def test_reg_and_dereg(self):
svc1 = ServiceInstance.builder().id("instance1").name("service1").build()
svc2 = ServiceInstance.builder().id("instance2").name("service1").build()
svc3 = ServiceInstance.builder().id("foo1").name("foo").build()
self.discovery.registerService(svc1)
self.discovery.registerService(svc2)
self.discovery.registerService(svc3)
self.assertTrue(self.client.exists(self.discovery.pathForInstance(svc1.getName(),svc1.getId())))
self.assertTrue(self.client.exists(self.discovery.pathForInstance(svc2.getName(),svc2.getId())))
self.assertTrue(self.client.exists(self.discovery.pathForInstance(svc3.getName(),svc3.getId())))
self.discovery.unregisterService(svc1)
self.discovery.unregisterService(svc2)
self.discovery.unregisterService(svc3)
self.assertFalse(self.client.exists(self.discovery.pathForInstance(svc1.getName(),svc1.getId())))
self.assertFalse(self.client.exists(self.discovery.pathForInstance(svc2.getName(),svc2.getId())))
self.assertFalse(self.client.exists(self.discovery.pathForInstance(svc3.getName(),svc3.getId())))
def test_query(self):
svc1 = ServiceInstance.builder().id("instance1").name("service1").build()
svc2 = ServiceInstance.builder().id("instance2").name("service1").build()
svc3 = ServiceInstance.builder().id("foo1").name("foo").build()
self.discovery.registerService(svc1)
self.discovery.registerService(svc2)
self.discovery.registerService(svc3)
self.assertEquals(sorted(["foo", "service1"]), sorted(self.discovery.queryForNames()))
instances = self.discovery.queryForInstances("service1")
self.assertEquals(2, len(instances))
self.assertEquals(sorted([svc1, svc2]), sorted(instances))
# make sure unregister works
self.discovery.unregisterService(svc2)
instances = self.discovery.queryForInstances("service1")
self.assertEquals(1, len(instances))
self.assertEquals(sorted([svc1]), sorted(instances))
instance = self.discovery.queryForInstance("service1","instance1")
self.assertTrue(instance)
self.assertEquals(svc1, instance)
class Server(threading.Thread):
'''
工作服务器（也是ZooKeeper的客户端）
'''
# 控制输出信息的锁，注意：这个是单机器的锁，这里实现的是分布式锁，并不存在本末倒置
print_mutex = threading.Lock()
DELAY_TIME = 3
def __init__(self, zk_server_address, lock_base_path, host, serve_mode):
threading.Thread.__init__(self)
# 锁的根节点路径
self.lock_base_path = lock_base_path
# 主机IP
self.host = host
# 工作模式，读/写
self.serve_mode = serve_mode
# 事件，初始化为False
self.event = threading.Event()
# 创建一个zookeeper客户端
self.zkclient = KazooClient(zk_server_address)
# 添加连接状态监听器
self.zkclient.add_listener(self.zk_connect_listener)
# 与zookeeper开启连接
self.zkclient.start()
# 连接状态监听器
def zk_connect_listener(self, state):
# 获取打印锁
Server.print_mutex.acquire()
if state == KeeperState.CONNECTED:
print self.host + " 已经开启..."
elif state == KazooState.LOST:
print self.host + " 停止服务..."
else:
raise Exception(self.host + " 未正常开启...")
# 获取打印锁
Server.print_mutex.release()
# 初始化
def run(self):
# 创建锁节点，形如/shared_lock/192.168.0.0-R-0000000001
self.create_lock_node()
# 获取锁
self.acquire_lock()
# 工作
self.work()
# 释放锁
self.release_lock()
# 准备停止
self.stop()
def create_lock_node(self):
# 先检查父节点，如果父节点不存在
if not self.zkclient.exists(self.lock_base_path):
# 先创建父节点
self.zkclient.create(self.lock_base_path)
# 拼凑出服务器子节点的完整路径
node_path = self.lock_base_path + "/" + self.host + "-" + self.serve_mode + "-"
# 创建临时顺序节点
self.node_path = self.zkclient.create(node_path, "", self.zkclient.default_acl, True, True)
# 删除事件的响应
def pre_node_delete_watch(self, data, stat, event):
if event and event.type == EventType.DELETED:
# 将事件设置为True
self.event.set()
# 获取锁
def acquire_lock(self):
# 提取出自己的节点名
node_name = self.node_path.split("/")[-1]
# 获取/shared_lock子节点排序列表
sorted_children = self.get_sorted_children()
# 得到节点的索引
node_index = sorted_children.index(node_name)
# 寻找最后一个写节点
def get_last_write_node_index():
# 逆向遍历
for i in range(node_index)[::-1]:
# 工作模式是节点名中的第二个部分
serve_mode = sorted_children[i].split("-")[1]
# 只要找到一个写请求，则立刻返回
if serve_mode == "W":
return i
# 如果全部都是读请求，则返回-1
return -1
# 如果是写请求，
if self.serve_mode == "W":
# 如果是，再判断自己是不是序号最小的节点
if node_index == 0:
# 立马返回，占用锁，开始写数据
return
# 如果不是，向比自己小的最后一个节点注册监听
else:
# 拼凑出前一个节点的路径
pre_node_path = self.lock_base_path + "/" + sorted_children[node_index - 1]
# 添加对前一个节点的删除事件的关注
self.zkclient.DataWatch(pre_node_path, self.pre_node_delete_watch)
# 这里应该等待锁
self.event.wait()
# 如果是读请求
else:
# 得到所有比自己小的子节点中的最后一个写节点的下标
last_write_node_index = get_last_write_node_index()
# 判断以下两个条件是否成立
# 1）没有比自己序号小的子节点
# 2）或是所有比自己小的子节点都是读请求
# 如果成立
if node_index == 0 or last_write_node_index < 0:
# 立马返回，占用共享锁，开始读数据
return
# 如果不成立，向比自己小的最后一个写节点注册监听
else:
# 拼凑出前一个节点的路径
pre_node_path = self.lock_base_path + "/" + sorted_children[last_write_node_index]
# 添加对前一个节点的删除事件的关注
self.zkclient.DataWatch(pre_node_path, self.pre_node_delete_watch)
# 这里应该等待锁
self.event.wait()
def work(self):
# 获取打印锁
Server.print_mutex.acquire()
# 如果是写请求，
if self.serve_mode == "W":
# 写一会数据，然后删除节点，关闭会话
print self.host + " 正在写数据..."
else:
# 读一会数据，然后删除节点，关闭会话
print self.host + " 正在读数据..."
Server.print_mutex.release()
# 这里暂停几秒钟。模拟工作耗时状态
sleep(self.DELAY_TIME)
# 释放锁
def release_lock(self):
# 删除自己的节点
self.zkclient.delete(self.node_path)
# 获取/shared_lock子节点排序列表
def get_sorted_children(self):
# 获取/shared_lock子节点列表
children = self.zkclient.get_children(self.lock_base_path)
###############################################################
# 这里sort函数的比较表达式是由两个函数实现，还挺有技巧的
###############################################################
# 返回节点的序列号
def get_lock_node_seq(node_name):
# 分割字符串，然后返回列表最后一个元素，先将其转化为整型
return string.atoi(node_name.split("-")[-1])
# 编号比较r函数
def sequence_compare(node1, node2):
return get_lock_node_seq(node1) - get_lock_node_seq(node2)
# 将列表排序
children.sort(cmp = sequence_compare)
return children
# 停止工作
def stop(self):
# 移除事件监听器
self.zkclient.remove_listener(self.pre_node_delete_watch)
# 会话
self.zkclient.stop()
self.zkclient.close()