# -*- coding: utf-8 -*- import functools import logging import threading import weakref import environs from time import time from collections import defaultdict import io from petrel_client.common import mem_trace from petrel_client.common.exception import ObjectNotFoundError LOG = logging.getLogger(__name__) ENV = environs.Env() class StatItem(object): __slots__ = ['op_name', 'total_io', 'total_hit', 'total_time', 'total_error', 'total_miss', 'error_count', 'total_byte' ] def __init__(self, op_name): self.op_name = op_name self.reset() def reset(self): self.total_io = 0 self.total_hit = 0 self.total_time = 0.0 self.total_error = 0 self.total_miss = 0 self.total_byte = 0 self.error_count = defaultdict(lambda: 0) @property def time_avg(self): return self.total_time / self.total_io if self.total_io else .0 @property def hit_ratio(self): return 1.0 * self.total_hit / self.total_io if self.total_io else .0 @property def speed(self): return 1.0 * self.total_byte / self.total_time if self.total_time else .0 def stat_io(self, callback=None): stat_info = f'{self.op_name} [total: {self.total_io}' \ f', hit: {self.total_hit}' \ f', miss: {self.total_miss}' \ f', error: {self.total_error}' \ f', time: {self.total_time:.6} s' \ f', time_avg: {self.time_avg:.6} s' \ f', hit ratio: {self.hit_ratio:.2%}' \ f', bytes: {_sizeof_fmt(self.total_byte)}' \ f', speed: {_sizeof_fmt(self.speed,suffix="B/s")}' \ f']' if self.error_count: items = ["{}: {}".format(k, v) for (k, v) in self.error_count.items()] stat_info = f'{stat_info}, error_count: [{", ".join(items)}]' if callback: callback(stat_info) else: LOG.info(stat_info) self.reset() class StatItemDict(dict): def __missing__(self, key): item = self[key] = StatItem(key) return item def stat_io(self, callback=None): for item in self.values(): item.stat_io(callback) class ClientStat(object): def __init__(self, client_id, name): self.client_id = client_id self.name = name self.stat_item_dict = StatItemDict() profiler = Profiler.get() self.profiler = profiler profiler.register(self) def __getitem__(self, op_name): return self.stat_item_dict[op_name] # 若使用 multiprocessing-logging，进程退出时候调用 __del__ 还存在问题 # def __del__(self): # # 这里有可能是再 python 将要退出的时候触发，此时file log已经不存在，会发生异常 # try: # self.profiler.unregister(self) # if self.total_io: # self.stat_io() # except Exception: # pass @property def total_io(self): return sum([item.total_io for item in self.stat_item_dict.values()]) @property def get_hit(self): return sum([item.total_hit for item in self.stat_item_dict.values() if item.op_name == 'get']) def stat_io(self, callback=None): stat_item_info_list = [] def cb(info): stat_item_info_list.append(info) for stat_item in self.stat_item_dict.values(): stat_item.stat_io(cb) if stat_item_info_list: stat_itme_info = ', '.join(stat_item_info_list) else: stat_itme_info = 'No IO operations' stat_info = '{}: {}'.format(self.name, stat_itme_info) if callback: callback(stat_info) else: LOG.info(stat_info) def profile(op_name): assert isinstance(op_name, str) def wrap(fn): @functools.wraps(fn) def new_fn(self, *args, **kwargs): return _profile(op_name, fn, self, *args, **kwargs) return new_fn return wrap def _profile(op_name, fn, client, *args, **kwargs): stat: StatItem = client.client_stat[op_name] start = time() try: ret = fn(client, *args, **kwargs) if isinstance(ret, (tuple, list)): content = ret[0] else: content = ret if isinstance(content, bytes): stat.total_byte += len(content) elif isinstance(content, int): stat.total_byte += content elif hasattr(content, 'content_length'): stat.total_byte += content.content_length elif op_name == 'get' and content is None: raise ObjectNotFoundError() stat.total_hit += 1 return ret except ObjectNotFoundError: stat.total_miss += 1 raise except Exception as e: stat.total_error += 1 err_name = e.__class__.__name__ stat.error_count[err_name] += 1 raise finally: end = time() stat.total_time += (end - start) stat.total_io += 1 client.client_stat.profiler.inc_op_count() class Profiler(object): thread_local = threading.local() default_conf = None @staticmethod def set_default_conf(conf): Profiler.default_conf = conf @staticmethod def get(): profiler = getattr(Profiler.thread_local, 'profiler', None) if not profiler: profiler = Profiler(Profiler.default_conf) setattr(Profiler.thread_local, 'profiler', profiler) return profiler def __init__(self, conf, *args, **kwargs): assert conf is not None self.stat_dict = weakref.WeakValueDictionary() self.op_count = 0 self.count_disp = ENV.int( 'count_disp', None) or conf.get_int('count_disp') self.enable_mem_trace = conf.get_boolean('enable_mem_trace') if self.enable_mem_trace: mem_trace.start() def register(self, client_stat: ClientStat): client_id = client_stat.client_id self.stat_dict[client_id] = client_stat def unregister(self, client_stat: ClientStat): client_id = client_stat.client_id del self.stat_dict[client_id] def inc_op_count(self): self.op_count += 1 if self.count_disp: if self.op_count >= self.count_disp: self.stat_io() self.op_count = 0 @staticmethod def set_count_disp(count_disp): if count_disp < 0: LOG.error('count_disp must be a nonnegative integer, actual value: %s', count_disp) return profiler = Profiler.get() profiler.count_disp = count_disp def stat_io(self): if LOG.isEnabledFor(logging.INFO): io_dict = { client_stat.name: client_stat.get_hit for client_stat in self.stat_dict.values()} total_io = sum(io_dict.values()) or 1 percentage = [f'{client_name}: {1.0 * count / total_io :.2%}' for client_name, count in io_dict.items()] for client_stat in self.stat_dict.values(): client_stat.stat_io() LOG.info('IO Percentage: %s', ', '.join(percentage)) if self.enable_mem_trace: snapshot = mem_trace.take_snapshot() buffer = io.StringIO() snapshot.display_top(buffer=buffer) LOG.info('Memory trace: \n%s', buffer.getvalue()) def enable(self): raise NotImplementedError() def disable(self): raise NotImplementedError() def _sizeof_fmt(num, suffix='B'): for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']: if abs(num) < 1024.0: return "%3.1f%s%s" % (num, unit, suffix) num /= 1024.0 return "%.1f%s%s" % (num, 'Yi', suffix)