import json import os from datetime import datetime, timedelta, timezone from uuid import uuid4 import boto3 import tuotuo from sinks import FullActivityClient from tuotuo.activities import Activity, ActivityClient, ActivityType, Sevirity 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"] def dict_to_s3(dict, bucket, key): s3 = boto3.resource("s3") res = s3.Bucket(bucket).put_object(Key=key, Body=json.dumps(dict)) return res async def write_activity_to_hub(activity: Activity): print(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_oom_detect_time(activity_id): print(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 int(activity.detect_time) def get_oom_evidence(activity_id): client = boto3.client("s3") response = client.get_object(Bucket=BUCKET, Key=f"oom_detector/{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 """ oom_client = tuotuo.client("oom") running_sqls = None 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() 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() running_sqls = entity running_sqls = { "topmem": running_sqls.top_memory, "toptime": running_sqls.top_time, } 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 - 3600, end=end, ) for activity in activities[:]: if activity.activity_name != "OOM Risk": activities.remove(activity) 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 async def get_evidences_from_tidb( cluster_id: str, instance: str, start_ts: datetime, end_ts: datetime ): """ input: cluster_id, start, end output: expensive_sqls, out of memory quota, memory exceeds quota """ client = tuotuo.client("logs") expensive_sqls = [] out_of_memory_quota = [] memory_exceeds_quota = [] async for entity in client.read( int(start_ts), int(end_ts), cluster_ids=[int(cluster_id)], components=["tidb"], ): if "Out Of Memory Quota" in entity.content and instance == entity.instance: out_of_memory_quota.append(entity.content) elif "memory exceeds quota" in entity.content and instance == entity.instance: memory_exceeds_quota.append(entity.content) oom_client = tuotuo.client("oom") async for entity in oom_client.read_expensive_query( int(start_ts), int(end_ts), cluster_ids=[int(cluster_id)], ): if instance == entity.instance: expensive_sql = entity._asdict() expensive_sql["time"] = expensive_sql["time"].timestamp() expensive_sqls.append(expensive_sql) return { "expensive_sqls": expensive_sqls, "out_of_memory_quota": out_of_memory_quota, "memory_exceeds_quota": memory_exceeds_quota, } async def oom_analyzer(oom_event_id, cluster_id, instance, occurrence_ts, tenant_id, project_id): """ output: summary running sql expensive sql other_clue out of memory quota memory exceeds quota """ analyze_time = int(datetime.now().timestamp()) activity_id = str(uuid4()) end_ts = occurrence_ts start_ts = end_ts - 0.5 * 3600 latest_oom_risk_event = await find_latest_oom_risk_event(cluster_id, instance, end_ts) if latest_oom_risk_event is not None: if latest_oom_risk_event.occurrence_time > start_ts: start_ts = latest_oom_risk_event.occurrence_time print(f"latest oom risk time {start_ts}, there is {end_ts - start_ts} seconds...") evidences = await get_evidences_from_tidb(cluster_id, instance, start_ts, end_ts) running_sqls = await get_running_sqls(cluster_id, instance, start_ts, end_ts) detect_ts = await get_oom_detect_time(oom_event_id) summary = f"At {datetime.fromtimestamp(occurrence_ts,tz=timezone.utc)}, cluster {cluster_id} - instance {instance} has been restarted due to OOM." details = { "cluster_id": cluster_id, "instance": instance, "occurrence_time": occurrence_ts, "detect_time": detect_ts, "analyze_time": analyze_time, "evidences": get_oom_evidence(oom_event_id), "running_sqls_topmem": running_sqls["topmem"], "running_sqls_toptime": running_sqls["toptime"], "expensive_sqls": evidences["expensive_sqls"], "logs": { "out_of_memory_quota": evidences["out_of_memory_quota"], "memory_exceeds_quota": evidences["memory_exceeds_quota"], }, } res = dict_to_s3(details, BUCKET, f"{PATH}/{activity_id}.json") payload = { "summary": summary, "cluster_id": cluster_id, "instance": instance, "occurrence_time": occurrence_ts, "detect_time": detect_ts, "analyze_time": analyze_time, "details": { "type": "s3", "bucket": BUCKET, "key": f"{PATH}/{activity_id}.json", }, } oom_analysis_action = Activity( activity_type=ActivityType.ACTION, activity_name="OOM Analysis", target_id=cluster_id, calibration_time=analyze_time, source=SOURCE, catalog="OOM", severity=Sevirity.DEBUG, activity_id=activity_id, application_id=tenant_id, concept_id=project_id, author="oom-analyzer", context_id=oom_event_id, context=[latest_oom_risk_event.activity_id], payload=payload, source_key=SOURCE_KEY, ) event_res = await write_activity_to_hub(oom_analysis_action) print(event_res) if __name__ == "__main__": import asyncio import time """ input: cluster_id, start, end, tenant_id, project_id output: oom_events excute: write details to s3 and write oom_events to activityhub non-goal: alert, invoke consumer non-goals are implemented in agent """ oom_event_id = "3b61bcad-968d-4f8e-9dc5-150d288353a3" cluster_id = "1379661944629574082" tenant_id = "1372813089198981282" project_id = "1372813089196881283" instance = "db-tidb-0" occurrence_ts = 1666791974 proc_start = time.perf_counter() oom_events = asyncio.run( oom_analyzer( oom_event_id, cluster_id, instance, occurrence_ts, tenant_id, project_id, ) ) proc_end = time.perf_counter() print(f"Elapsed time: {proc_end - proc_start} seconds")