class LoadBalanceChecker(Checker):
"""check loadbalance operations in a dependent thread"""
def __init__(self, collection_name=None):
super().__init__(collection_name=collection_name)
self.utility_wrap = ApiUtilityWrapper()
self.c_wrap.load(enable_traceback=enable_traceback)
def keep_running(self):
while True:
c_name = self.c_wrap.name
res, _ = self.c_wrap.get_replicas()
# prepare load balance params
# find a group which has multi nodes
group_nodes = []
for g in res.groups:
if len(g.group_nodes) >= 2:
group_nodes = list(g.group_nodes)
break
src_node_id = group_nodes[0]
dst_node_ids = group_nodes[1:]
res, _ = self.utility_wrap.get_query_segment_info(c_name)
segment_distribution = cf.get_segment_distribution(res)
sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
# load balance
t0 = time.time()
_, result = self.utility_wrap.load_balance(c_name, src_node_id,
dst_node_ids,
sealed_segment_ids)
t1 = time.time()
# get segments distribution after load balance
time.sleep(3)
res, _ = self.utility_wrap.get_query_segment_info(c_name)
segment_distribution = cf.get_segment_distribution(res)
sealed_segment_ids_after_load_banalce = segment_distribution[
src_node_id]["sealed"]
check_1 = len(
set(sealed_segment_ids)
& set(sealed_segment_ids_after_load_banalce)) == 0
des_sealed_segment_ids = []
for des_node_id in dst_node_ids:
des_sealed_segment_ids += segment_distribution[des_node_id][
"sealed"]
# assert sealed_segment_ids is subset of des_sealed_segment_ids
check_2 = set(sealed_segment_ids).issubset(
set(des_sealed_segment_ids))
if result and (check_1 and check_2):
self.rsp_times.append(t1 - t0)
self.average_time = (
(t1 - t0) + self.average_time * self._succ) / (self._succ +
1)
self._succ += 1
log.debug(
f"load balance success, time: {t1 - t0:.4f}, average_time: {self.average_time:.4f}"
)
else:
self._fail += 1
sleep(10)
class LoadBalanceChecker(Checker):
"""check loadbalance operations in a dependent thread"""
def __init__(self, collection_name=None):
if collection_name is None:
collection_name = cf.gen_unique_str("LoadBalanceChecker_")
super().__init__(collection_name=collection_name)
self.utility_wrap = ApiUtilityWrapper()
self.c_wrap.load()
self.sealed_segment_ids = None
self.dst_node_ids = None
self.src_node_id = None
@trace()
def load_balance(self):
res, result = self.utility_wrap.load_balance(
self.c_wrap.name, self.src_node_id, self.dst_node_ids, self.sealed_segment_ids)
return res, result
def prepare(self):
"""prepare load balance params"""
res, _ = self.c_wrap.get_replicas()
# find a group which has multi nodes
group_nodes = []
for g in res.groups:
if len(g.group_nodes) >= 2:
group_nodes = list(g.group_nodes)
break
self.src_node_id = group_nodes[0]
self.dst_node_ids = group_nodes[1:]
res, _ = self.utility_wrap.get_query_segment_info(self.c_wrap.name)
segment_distribution = cf.get_segment_distribution(res)
self.sealed_segment_ids = segment_distribution[self.src_node_id]["sealed"]
@exception_handler()
def run_task(self):
self.prepare()
res, result = self.load_balance()
return res, result
def keep_running(self):
while self._keep_running:
self.run_task()
sleep(constants.WAIT_PER_OP / 10)