"""HTTP service for TiDB Cloud SQL Health Check inspections.""" import logging import os import queue import threading import time import traceback import uuid from typing import Any, Dict from flask import Flask, jsonify, request from service import handler as tablefullscan_handler from service import slowlog_top_digest_handler from service import tikv_high_cpu_handler from service import tikv_high_duration_handler logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s", ) logger = logging.getLogger(__name__) app = Flask(__name__) def _positive_int_from_env(name: str, default: int, fallback_name: str | None = None) -> int: raw_value = os.getenv(name) if raw_value is None and fallback_name: raw_value = os.getenv(fallback_name) value = int(raw_value if raw_value is not None else str(default)) if value <= 0: raise ValueError(f"{name} must be positive") return value MAX_CONCURRENT_INSPECTIONS = _positive_int_from_env( "SQL_HEALTHCHECK_MAX_CONCURRENT_INSPECTIONS", 2, fallback_name="TABLEFULLSCAN_MAX_CONCURRENT_INSPECTIONS", ) MAX_PENDING_INSPECTIONS = _positive_int_from_env( "SQL_HEALTHCHECK_MAX_PENDING_INSPECTIONS", 32, fallback_name="TABLEFULLSCAN_MAX_PENDING_INSPECTIONS", ) INSPECTION_QUEUE: queue.Queue[Dict[str, Any]] = queue.Queue(maxsize=MAX_PENDING_INSPECTIONS) QUEUE_STATE_LOCK = threading.Lock() ACTIVE_INSPECTIONS = 0 INTERNAL_ITEM_FIELD = "_sql_healthcheck_item" INSPECTION_HANDLERS = { "tablefullscan": { "name": tablefullscan_handler.ITEM_NAME, "validate": tablefullscan_handler.validate_request_params, "run": tablefullscan_handler.run_inspection, "build_s3_key": tablefullscan_handler.build_s3_key, }, "slowlog_top_digest": { "name": slowlog_top_digest_handler.ITEM_NAME, "validate": slowlog_top_digest_handler.validate_request_params, "run": slowlog_top_digest_handler.run_inspection, "build_s3_key": slowlog_top_digest_handler.build_s3_key, }, "tikv_high_duration": { "name": tikv_high_duration_handler.ITEM_NAME, "validate": tikv_high_duration_handler.validate_request_params, "run": tikv_high_duration_handler.run_inspection, "build_s3_key": tikv_high_duration_handler.build_s3_key, }, "tikv_high_cpu": { "name": tikv_high_cpu_handler.ITEM_NAME, "validate": tikv_high_cpu_handler.validate_request_params, "run": tikv_high_cpu_handler.run_inspection, "build_s3_key": tikv_high_cpu_handler.build_s3_key, }, } def _worker_loop(worker_id: int): global ACTIVE_INSPECTIONS while True: body = INSPECTION_QUEUE.get() item_key = body.get(INTERNAL_ITEM_FIELD, "tablefullscan") handler = INSPECTION_HANDLERS.get(item_key, INSPECTION_HANDLERS["tablefullscan"]) cluster_id = body.get("clusterId", "") report_id = body.get("reportId", "") with QUEUE_STATE_LOCK: ACTIVE_INSPECTIONS += 1 active = ACTIVE_INSPECTIONS queued = INSPECTION_QUEUE.qsize() logger.info( "Worker %d started inspection item=%s cluster=%s report_id=%s active=%d queued=%d", worker_id, handler["name"], cluster_id, report_id, active, queued, ) started_at = time.monotonic() try: request_body = dict(body) request_body.pop(INTERNAL_ITEM_FIELD, None) handler["run"](request_body) except Exception as exc: logger.error( "Queued inspection failed item=%s cluster=%s report_id=%s error=%s\n%s", handler["name"], cluster_id, report_id, exc, traceback.format_exc(), ) finally: INSPECTION_QUEUE.task_done() with QUEUE_STATE_LOCK: ACTIVE_INSPECTIONS -= 1 active = ACTIVE_INSPECTIONS queued = INSPECTION_QUEUE.qsize() logger.info( "Worker %d finished inspection item=%s cluster=%s report_id=%s active=%d queued=%d duration_s=%.3f", worker_id, handler["name"], cluster_id, report_id, active, queued, time.monotonic() - started_at, ) for worker_index in range(MAX_CONCURRENT_INSPECTIONS): worker = threading.Thread( target=_worker_loop, args=(worker_index + 1,), daemon=True, name=f"sql-health-check-worker-{worker_index + 1}", ) worker.start() @app.route("/healthz", methods=["GET"]) def healthz(): return jsonify({"status": "ok"}), 200 @app.route("/clinic-internal/healthcheck/sql/tablefullscan-slowquery", methods=["POST"]) @app.route("/api/v1/inspect", methods=["POST"]) def inspect_tablefullscan(): return _inspect("tablefullscan") @app.route("/clinic-internal/healthcheck/sql/slowlog-top-digest", methods=["POST"]) @app.route("/api/v1/inspect/slowlog-top-digest", methods=["POST"]) def inspect_slowlog_top_digest(): return _inspect("slowlog_top_digest") @app.route("/clinic-internal/healthcheck/sql/tikv-high-duration", methods=["POST"]) @app.route("/api/v1/inspect/tikv-high-duration", methods=["POST"]) def inspect_tikv_high_duration(): return _inspect("tikv_high_duration") @app.route("/clinic-internal/healthcheck/sql/tikv-high-cpu", methods=["POST"]) @app.route("/api/v1/inspect/tikv-high-cpu", methods=["POST"]) def inspect_tikv_high_cpu(): return _inspect("tikv_high_cpu") def _inspect(item_key: str): handler = INSPECTION_HANDLERS[item_key] try: body = request.get_json(force=True) except Exception: return jsonify({"error": "Invalid JSON body"}), 400 if not isinstance(body, dict): return jsonify({"error": "JSON body must be an object"}), 400 required = ["clusterId", "timeRange"] missing = [field for field in required if field not in body] if missing: return jsonify({"error": f"Missing required fields: {missing}"}), 400 time_range = body["timeRange"] if not isinstance(time_range, dict) or "start" not in time_range or "end" not in time_range: return jsonify({"error": "timeRange must contain 'start' and 'end'"}), 400 try: handler["validate"](body) except ValueError as exc: return jsonify({"error": str(exc)}), 400 cluster_id = body["clusterId"] report_id = body.get("reportId", str(uuid.uuid4())) body["reportId"] = report_id body[INTERNAL_ITEM_FIELD] = item_key s3_key = handler["build_s3_key"](cluster_id, report_id) s3_url = f"s3://delicated-healthcheck/{s3_key}" try: INSPECTION_QUEUE.put_nowait(body) except queue.Full: with QUEUE_STATE_LOCK: active = ACTIVE_INSPECTIONS queued = INSPECTION_QUEUE.qsize() logger.warning( "Rejecting inspection item=%s cluster=%s report_id=%s active=%d queued=%d max_concurrent=%d max_pending=%d", handler["name"], cluster_id, report_id, active, queued, MAX_CONCURRENT_INSPECTIONS, MAX_PENDING_INSPECTIONS, ) return jsonify({ "name": handler["name"], "status": "rejected", "error": "inspection queue is full", "clusterId": cluster_id, "reportId": report_id, }), 429 with QUEUE_STATE_LOCK: active = ACTIVE_INSPECTIONS queued = INSPECTION_QUEUE.qsize() logger.info( "Queued inspection item=%s cluster=%s report_id=%s active=%d queued=%d max_concurrent=%d max_pending=%d", handler["name"], cluster_id, report_id, active, queued, MAX_CONCURRENT_INSPECTIONS, MAX_PENDING_INSPECTIONS, ) return jsonify({ "name": handler["name"], "status": "processing", "s3_url": s3_url, "clusterId": cluster_id, "reportId": report_id, }), 202