from __future__ import annotations import base64 import json import os import pickle from collections import namedtuple from datetime import datetime, timezone from typing import Type from uuid import uuid4 import aiohttp import boto3 import tuotuo from sinks import FullActivityClient from task import BaseTask from tuotuo.activities import Activity, ActivityClient, ActivityType, Sevirity from tuotuo.logs.utils import Parser ACTIVITY_HUB_ENDPOINT = os.environ["ACTIVITY_HUB_ENDPOINT"] BUCKET = os.environ["ACTIVITY_BUCKET"] PATH = "oom_analyzer" SOURCE_KEY = os.environ["SOURCE_KEY"] SOURCE = os.environ["SOURCE"] ANALYZE_INTERVAL = int(os.environ["ANALYZE_SECONDS"]) SQLStats = namedtuple("SQLStats", ["time", "cluster", "component", "instance", "content"]) def dict_to_s3(dict, bucket, key): s3 = boto3.resource("s3") res = s3.Bucket(bucket).put_object(Key=key, Body=json.dumps(dict)) print(f"write to s3 {bucket}/{key}") return res async def get_latest_oom_record( cluster_id, instance, org_id, project_id, provider, region_name, start_ts, end_ts ): clinic_token = os.environ["CLINIC_API_TOKEN"] params = { "instances": instance, "start": start_ts, "end": end_ts, } timeout = aiohttp.ClientTimeout(total=10) print(f"start to get latest oom record: {params}") async with aiohttp.ClientSession(trust_env=True, timeout=timeout) as session: async with session.get( f"https://clinic.pingcap.com/ngm/api/v1/oom_record/list", params=params, headers={ "Authorization": f"Bearer {clinic_token}", "X-Provider": provider, "X-Region": region_name, "X-Org-Id": org_id, "X-Project-Id": project_id, "X-Cluster-Id": cluster_id, }, ) as rsp: if not rsp.ok: raise Exception(f"HTTP Status Code {rsp.status} ({await rsp.text()})") rsp = await rsp.json() last_record = None for item in rsp["records"]: if item["instance"] != instance.removeprefix("db-"): continue for file_meta in item["files"]: if "heap" not in file_meta["filename"]: continue if ( last_record is None or file_meta["created_at_unix_sec"] > last_record["created_at_unix_sec"] ): last_record = file_meta print(f"succeed to get latest oom record: {last_record}") return last_record async def get_top10_memory_usage_heap(provider, region_name, file_key): clinic_token = os.environ["CLINIC_API_TOKEN"] params = { "provider": provider, "region": region_name, "key": file_key, } timeout = aiohttp.ClientTimeout(total=30) print(f"start to get top10 memory usage heap: {params}") async with aiohttp.ClientSession(trust_env=True, timeout=timeout) as session: async with session.get( f"https://clinic.pingcap.com/ngm/api/v1/oom_record/top10_mem_usage", params=params, headers={"Authorization": f"Bearer {clinic_token}"}, ) as rsp: if not rsp.ok: raise Exception(f"HTTP Status Code {rsp.status} ({await rsp.text()})") rsp = await rsp.json() print(f"succeed to get top10 memory usage heap: {rsp}") return rsp async def write_activity_to_hub(activity: Activity): async with FullActivityClient( source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT ) as client: return await client.write_activity(activity) async def get_activity(activity_id): async with ActivityClient( source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT ) as client: activity = await client.read_activity(activity_id) return activity def get_evidence(activity_id, activity_name): client = boto3.client("s3") path = activity_name.replace(" ", "_") + "_detector" response = client.get_object(Bucket=BUCKET, Key=f"{path}/{activity_id}.json") response = json.loads(response["Body"].read().decode("utf-8")) return response["evidence"] async def get_running_sqls(cluster_id, instance, start_ts, end_ts): """ input: cluster_id, start, end output: running_sqls, 2 lists """ print( f"start to get running sqls, cluster_id: {cluster_id}, instance: {instance}, start_ts: {start_ts}, end_ts: {end_ts}" ) oom_client = tuotuo.client("oom") running_sqls = {"topmem": [], "toptime": []} async for entity in oom_client.read_running_sql( int(start_ts), int(end_ts), cluster_ids=[int(cluster_id)], ): if instance == "db-" + entity.instance: for i, item in enumerate(entity.top_memory[:]): entity.top_memory[i] = item._asdict() entity.top_memory[i]["time"] = entity.top_memory[i]["time"].timestamp() if entity.top_memory[i]["content"].get("mem_max"): mem_max_human = entity.top_memory[i]["content"]["mem_max"][ entity.top_memory[i]["content"]["mem_max"].find("(") + 1 : entity.top_memory[i]["content"]["mem_max"].find(")") ] else: mem_max_human = None entity.top_memory[i]["content"]["mem_max_human"] = mem_max_human for i, item in enumerate(entity.top_time[:]): entity.top_time[i] = item._asdict() entity.top_time[i]["time"] = entity.top_time[i]["time"].timestamp() if entity.top_time[i]["content"].get("mem_max"): mem_max_human = entity.top_time[i]["content"]["mem_max"][ entity.top_time[i]["content"]["mem_max"].find("(") + 1 : entity.top_time[i]["content"]["mem_max"].find(")") ] else: mem_max_human = None entity.top_time[i]["content"]["mem_max_human"] = mem_max_human res = entity running_sqls = {"topmem": res.top_memory, "toptime": res.top_time} print(f"succeed to get running sqls.") return running_sqls async def find_latest_oom_risk_event(cluster_id, instance, end): async with ActivityClient( source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT ) as client: activities = await client.search_activities( target_id=cluster_id, start=end - ANALYZE_INTERVAL, end=end, ) print("start: ", end - ANALYZE_INTERVAL, "end: ", end) for activity in activities[:]: if activity.activity_name != "OOM Risk": activities.remove(activity) continue if activity.payload.get("instance") != instance: activities.remove(activity) if len(activities) == 0: return None else: latest = activities[0] for activity in activities: if activity.occurrence_time > latest.occurrence_time: latest = activity return latest class OOMAnalyzeTask(BaseTask): activity_id: str cluster_id: str instance: str tenant_id: str project_id: str occurrence_ts: int provider: str region_name: str def __init__( self, activity_id: str, cluster_id: str, instance: str, tenant_id: str, project_id: str, occurrence_ts: int, clinic_org_id: str, provider: str, region_name: str, **kwargs, ) -> None: self.related_event_id = activity_id self.cluster_id = cluster_id self.instance = instance self.tenant_id = tenant_id self.project_id = project_id self.occurrence_ts = occurrence_ts self.clinic_org_id = clinic_org_id self.provider = provider self.region_name = region_name self.analyze_time = None self.activity_id = None self.start_ts = None self.latest_oom_risk_event = None self.detect_ts = None self.related_activity = None self.summary = None self.start_ts = None self.related_evidence = None self.running_sqls = None self.top10_memory_usage_heap = None self.risk_evidence = { "expensive_sqls": [], "out_of_memory_quota": [], "memory_exceeds_quota": [], "global_memory_controller": [], } self.ctx_mark = None self.expensive_query_status = False self.other_clue_status = False self.reset_start_ts = None self.retry = 0 @classmethod def resume(cls, ctx: str) -> OOMAnalyzeTask: print(f"resume from context {ctx}") s3_object = ctx s3 = boto3.client("s3") tmpfile = "/tmp/resume.pkl" with open(tmpfile, "wb") as f: s3.download_fileobj(BUCKET, s3_object, f) task = pickle.load(open(tmpfile, "rb")) task.start_ts = int(task.ctx_mark) task.retry += 1 return task async def run(self) -> dict: # get gcp key try: s3 = boto3.client("s3") key = s3.get_object(Bucket=BUCKET, Key="secrets/gcp_secret")["Body"] os.environ["GCP_PRIVATE_KEY"] = bytes.decode(base64.b64decode(key.read())) except Exception as e: print(f"get gcp key failed: {e}") if self.retry > 6: print("retry over 90 minutes, give up") return { "status": "failed", "reason": "retry over 90 minutes, give up", "info": { "cluster_id": self.cluster_id, "instance": self.instance, "occurrence_ts": self.occurrence_ts, "related_event_id": self.related_event_id, }, } if self.analyze_time is None: self.analyze_time = int(datetime.now().timestamp()) if self.activity_id is None: self.activity_id = str(uuid4()) if self.related_activity is None: self.related_activity = await get_activity(self.related_event_id) if self.detect_ts is None: self.detect_ts = int(self.related_activity.detect_time) if self.summary is None: if self.related_activity.activity_name == "tidb oom": self.summary = f"At {datetime.fromtimestamp(self.occurrence_ts,tz=timezone.utc)}, cluster {self.cluster_id} - instance {self.instance} was restarted due to OOM." elif self.related_activity.activity_name == "oom risk": self.summary = f"At {datetime.fromtimestamp(self.occurrence_ts,tz=timezone.utc)}, cluster {self.cluster_id} - instance {self.instance} encountered an OOM risk." if self.latest_oom_risk_event is None: if self.related_activity.activity_name == "tidb oom": self.latest_oom_risk_event = await find_latest_oom_risk_event( self.cluster_id, self.instance, self.occurrence_ts ) elif self.related_activity.activity_name == "oom risk": self.latest_oom_risk_event = self.related_activity if self.start_ts is None: self.start_ts = self.occurrence_ts - ANALYZE_INTERVAL if self.related_activity.activity_name == "tidb oom": if self.latest_oom_risk_event is not None: if self.latest_oom_risk_event.occurrence_time > self.start_ts: self.start_ts = self.latest_oom_risk_event.occurrence_time if self.reset_start_ts is None: self.reset_start_ts = self.start_ts if self.related_evidence is None: self.related_evidence = get_evidence( self.related_event_id, self.related_activity.activity_name ) if self.running_sqls is None: # the condition of "oom risk" log and dump files are the same, but there is an interval between them. # log is generated befor dumping files, the minimum interval of two pieces of the same log is 10s. # so we use a 10s interval to get the running sqls. if self.related_activity.activity_name == "tidb oom": self.running_sqls = await get_running_sqls( self.cluster_id, self.instance, self.start_ts, self.occurrence_ts, ) elif self.related_activity.activity_name == "oom risk": self.running_sqls = await get_running_sqls( self.cluster_id, self.instance, self.occurrence_ts, self.occurrence_ts + 10, ) if self.top10_memory_usage_heap is None: last_record = await get_latest_oom_record( self.cluster_id, self.instance, self.tenant_id, self.project_id, self.provider, self.region_name, self.start_ts, self.occurrence_ts, ) if last_record is None: self.top10_memory_usage_heap = {} else: self.top10_memory_usage_heap = await get_top10_memory_usage_heap( self.provider, self.region_name, last_record["key"] ) if self.expensive_query_status == False: log_client = tuotuo.client("logs") print( f"start to get expensive query logs, cluster_id: {self.cluster_id}, instance: {self.instance}, start_ts: {self.start_ts}, end_ts: {self.occurrence_ts}" ) with Parser() as p: async for entity in log_client.read( int(self.start_ts), int(self.occurrence_ts), cluster_ids=[int(self.cluster_id)], components=["tidb"], ): self.ctx_mark = entity.time.timestamp() if ( ( len(self.risk_evidence["out_of_memory_quota"]) > 0 and self.ctx_mark < self.risk_evidence["out_of_memory_quota"][-1]["time"] ) or ( len(self.risk_evidence["memory_exceeds_quota"]) > 0 and self.ctx_mark < self.risk_evidence["memory_exceeds_quota"][-1]["time"] ) or ( len(self.risk_evidence["expensive_sqls"]) > 0 and self.ctx_mark < self.risk_evidence["expensive_sqls"][-1]["time"] ) or ( len(self.risk_evidence["global_memory_controller"]) > 0 and self.ctx_mark < self.risk_evidence["global_memory_controller"][-1]["time"] ) ): continue if "Out Of Memory Quota" in entity.content and self.instance == entity.instance: self.risk_evidence["out_of_memory_quota"].append( { "time": entity.time.timestamp(), "content": entity.content, } ) if ( "global memory controller" in entity.content and self.instance == entity.instance ): self.risk_evidence["global_memory_controller"].append( { "time": entity.time.timestamp(), "content": entity.content, } ) if ( "memory exceeds quota" in entity.content and self.instance == entity.instance and ( "txn_start_ts" not in entity.content or "mem_max" not in entity.content or "sql" not in entity.content ) ): self.risk_evidence["memory_exceeds_quota"].append( { "time": entity.time.timestamp(), "content": entity.content, } ) if not entity.content.startswith("["): continue if ( ("expensive_query" in entity.content) or ( "memory exceeds quota" in entity.content and "txn_start_ts" in entity.content and "mem_max" in entity.content and "sql" in entity.content ) ) and self.instance == entity.instance: x = p.parse(entity._replace(component="tidb")) if x.message == "expensive_query" or x.message == "memory exceeds quota": fields = dict([(f["Name"], f["Value"]) for f in x.fields]) expensive = SQLStats( entity.time, entity.cluster, entity.component, entity.instance, fields, ) if self.instance == expensive.instance: expensive_sql = expensive._asdict() expensive_sql["time"] = expensive_sql["time"].timestamp() if expensive_sql["content"].get("mem_max"): mem_max_human = expensive_sql["content"]["mem_max"][ expensive_sql["content"]["mem_max"].find("(") + 1 : expensive_sql["content"]["mem_max"].find(")") ] else: mem_max_human = None expensive_sql["content"]["mem_max_human"] = mem_max_human self.risk_evidence["expensive_sqls"].append(expensive_sql) print(f"succeed to get expensive query logs.") self.expensive_query_status = True if self.expensive_query_status == True: # and self.other_clue_status == True: clinic = tuotuo.client("clinic") cluster_meta_info = await clinic.get_cluster_meta(self.clinic_org_id, self.cluster_id) print(f"get cluster meta info: {cluster_meta_info}") self.details = { "cluster_id": self.cluster_id, "instance": self.instance, "occurrence_time": self.occurrence_ts, "detect_time": self.detect_ts, "analyze_time": self.analyze_time, "evidences": self.related_evidence, "running_sqls_topmem": self.running_sqls["topmem"], "running_sqls_toptime": self.running_sqls["toptime"], "expensive_sqls": self.risk_evidence["expensive_sqls"], "version": cluster_meta_info["version"], "components": cluster_meta_info["components"], "logs": { "out_of_memory_quota": self.risk_evidence["out_of_memory_quota"], "memory_exceeds_quota": self.risk_evidence["memory_exceeds_quota"], "global_memory_controller": self.risk_evidence["global_memory_controller"], }, "heap": self.top10_memory_usage_heap, } _ = dict_to_s3(self.details, BUCKET, f"{PATH}/{self.activity_id}.json") payload = { "summary": self.summary, "cluster_id": self.cluster_id, "instance": self.instance, "occurrence_time": self.occurrence_ts, "detect_time": self.detect_ts, "analyze_time": self.analyze_time, "details": { "type": "s3", "bucket": BUCKET, "key": f"{PATH}/{self.activity_id}.json", }, } if self.related_activity.activity_name == "tidb oom": context = ( [self.latest_oom_risk_event.activity_id] if self.latest_oom_risk_event is not None else [] ) else: context = [] oom_analysis_action = Activity( activity_type=ActivityType.ACTION, activity_name="OOM Analysis", target_id=self.cluster_id, calibration_time=self.analyze_time, source=SOURCE, catalog="OOM", severity=Sevirity.DEBUG, activity_id=self.activity_id, application_id=self.tenant_id, concept_id=self.project_id, author="oom-analyzer", context_id=self.related_event_id, context=context, payload=payload, source_key=SOURCE_KEY, ) print(f"start to write event to activity hub.") event_res = await write_activity_to_hub(oom_analysis_action) if event_res.get("status") == "ok": print( f"[Info][Write event Successful] {self.cluster_id} OOM analyzer action: {oom_analysis_action}" ) else: print( f"[Error][Write event Failed] {self.cluster_id} OOM analyzer action: {oom_analysis_action}" ) def suspend(self) -> str: if self.ctx_mark == None: raise TimeoutError("task timeout before checkpoint") s3_object = f"oom_analyzer/oom_analyzer-{self.activity_id}-{int(self.ctx_mark)}.pkl" s3 = boto3.client("s3") if not os.path.exists("/tmp/oom_analyzer"): os.makedirs("/tmp/oom_analyzer") with open(f"/tmp/{s3_object}", "wb") as handle: pickle.dump(self, handle) try: s3.upload_file(f"/tmp/{s3_object}", BUCKET, s3_object) except Exception as e: print(e) return False print(s3_object) return s3_object def task_class() -> Type[BaseTask]: return OOMAnalyzeTask