#!/usr/bin/env python3 """ manual_compaction.py - Trigger major compaction on cloud-storage-engine regions by scanning a key range via PD, discovering each region's leader, and issuing compaction requests to the leader store's status endpoint. This replaces the previous store-based approach (--store-id / /regions/store) with a range-based traversal that is safe for large clusters and restartable through checkpoints. Key format ---------- --start-key and --end-key accept lowercase hex-encoded byte strings, e.g. "41" for byte 0x41 ('A'). They are passed directly to PD's /pd/api/v1/regions/key endpoint with format=hex. PD returns start_key / end_key in each RegionInfo as **hex-encoded** strings. We decode them with bytes.fromhex() when advancing the scan cursor. Checkpoint ---------- If --checkpoint-file is specified, the script persists the current scan key and counters periodically (no more often than --checkpoint-interval seconds) and at completion. On restart it reads the checkpoint and resumes from the stored key. Dry-run ------- With --dry-run the script evaluates eligibility and logs what it *would* do, but does not POST any compaction request and does not write or update checkpoint files (safe read-only mode). Example ------- # Compact all regions from key 0x41 ('A') to key 0x5a ('Z'): python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 \\ --start-key 41 \\ --end-key 5a \\ --old-ratio 0.3 \\ --checkpoint-file /tmp/compact_checkpoint.json # Compact the full keyspace (empty start/end = beginning/end): python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 # Resume after interruption: python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 \\ --start-key 41 \\ --end-key 5a \\ --checkpoint-file /tmp/compact_checkpoint.json """ import argparse import json import logging import os import sys import tempfile import time import urllib.parse import urllib.request import urllib.error import ssl from datetime import datetime, timezone CHECKPOINT_VERSION = 1 PD_REGIONS_KEY_PATH = "/pd/api/v1/regions/key" PD_STORE_PATH = "/pd/api/v1/store" DEFAULT_REGION_LIMIT = 512 DEFAULT_OLD_RATIO = 0.3 DEFAULT_INTERVAL = 1 DEFAULT_TIMEOUT = 30 DEFAULT_CHECKPOINT_INTERVAL = 60 DEFAULT_SCAN_INTERVAL = 0.5 logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", datefmt="%Y-%m-%dT%H:%M:%S", ) logger = logging.getLogger(__name__) def make_ssl_context(args): if not args.ca_cert and not args.client_cert: return None ctx = ssl.create_default_context() if args.ca_cert: ctx.load_verify_locations(args.ca_cert) if args.client_cert and args.client_key: ctx.load_cert_chain(args.client_cert, args.client_key) return ctx def normalize_status_address(addr, use_tls): """ Normalize a raw status_address (host:port or host:port) to an HTTP/HTTPS URL. If the address already starts with http:// or https://, it is used as-is. """ if addr.startswith("http://") or addr.startswith("https://"): return addr scheme = "https" if use_tls else "http" return f"{scheme}://{addr}" def pd_request(pd_addrs, path, ssl_ctx, timeout, params=None): """ Try each PD address in order until one succeeds. Returns the decoded JSON response body. Raises RuntimeError if all addresses fail. """ query = ("?" + urllib.parse.urlencode(params)) if params else "" last_err = None for addr in pd_addrs: url = addr.rstrip("/") + path + query try: req = urllib.request.Request(url) handler = urllib.request.HTTPSHandler(context=ssl_ctx) if ssl_ctx else urllib.request.HTTPHandler() opener = urllib.request.build_opener(handler) with opener.open(req, timeout=timeout) as resp: body = resp.read() return json.loads(body) except Exception as exc: logger.debug("PD request failed for %s: %s", addr, exc) last_err = exc raise RuntimeError(f"All PD addresses failed for {path}: {last_err}") def http_post(url, ssl_ctx, timeout, params=None): """ POST to a TiKV status endpoint URL with optional query params. Returns (status_code, body_str). """ if params: url = url + "?" + urllib.parse.urlencode(params) try: req = urllib.request.Request(url, data=b"", method="POST") handler = urllib.request.HTTPSHandler(context=ssl_ctx) if ssl_ctx else urllib.request.HTTPHandler() opener = urllib.request.build_opener(handler) with opener.open(req, timeout=timeout) as resp: return resp.status, resp.read().decode("utf-8", errors="replace") except urllib.error.HTTPError as exc: return exc.code, exc.read().decode("utf-8", errors="replace") except Exception as exc: raise def http_get(url, ssl_ctx, timeout): """ GET a TiKV status endpoint URL. Returns the decoded JSON body or raises on error. """ try: req = urllib.request.Request(url) handler = urllib.request.HTTPSHandler(context=ssl_ctx) if ssl_ctx else urllib.request.HTTPHandler() opener = urllib.request.build_opener(handler) with opener.open(req, timeout=timeout) as resp: return json.loads(resp.read()) except urllib.error.HTTPError as exc: raise RuntimeError(f"HTTP {exc.code} from {url}: {exc.read().decode('utf-8', errors='replace')}") def scan_regions_batch(pd_addrs, start_key_raw, end_key_raw, limit, ssl_ctx, timeout): """ Call GET /pd/api/v1/regions/key with hex-encoded key bytes. PD's ScanRegions handler passes the key/end_key params through apiutil.ParseHexKeys; with format=hex the server hex-decodes the strings into raw bytes before scanning. Empty string means "beginning" (key) or "end of all keys" (end_key). Returns a list of region dicts from the PD JSON response. """ params = {} # Encode raw bytes as lowercase hex; pass format=hex so PD decodes them. params["key"] = start_key_raw.hex() if start_key_raw else "" if end_key_raw: params["end_key"] = end_key_raw.hex() params["limit"] = str(limit) params["format"] = "hex" data = pd_request(pd_addrs, PD_REGIONS_KEY_PATH, ssl_ctx, timeout, params=params) return data.get("regions") or [] def get_store_status_address(pd_addrs, store_id, ssl_ctx, timeout, use_tls, store_cache): """ Fetch store info from PD and return the normalized status address URL. Results are cached in store_cache keyed by store_id. Falls back to 'address' only when 'status_address' is absent. """ if store_id in store_cache: return store_cache[store_id] path = f"{PD_STORE_PATH}/{store_id}" try: data = pd_request(pd_addrs, path, ssl_ctx, timeout) except Exception as exc: raise RuntimeError(f"Failed to fetch store {store_id} from PD: {exc}") store_meta = data.get("store", {}) raw_addr = ( store_meta.get("status_address") or store_meta.get("status-address") or store_meta.get("address") ) if not raw_addr: raise RuntimeError(f"Store {store_id} has no usable address in PD response") url = normalize_status_address(raw_addr, use_tls) store_cache[store_id] = url return url def get_region_engine_stats(base_url, region_id, ssl_ctx, timeout): """ GET /kvengine/{region_id} from a TiKV status address. Returns a dict with 'active', 'entries', 'old_entries', and 'tombs', normalizing the status API's 'old-entries' field to 'old_entries'. """ url = f"{base_url}/kvengine/{region_id}" data = http_get(url, ssl_ctx, timeout) active = data.get("active", False) entries = data.get("entries", 0) old_entries = data.get("old-entries", 0) tombs = data.get("tombs", data.get("tombstones", 0)) return { "active": active, "entries": entries, "old_entries": old_entries, "tombs": tombs, } def is_eligible(stats, old_ratio_threshold): """ Return True if the region should be compacted: - region is active - entries + old_entries > 0 - (tombs + old_entries) / (entries + old_entries) > threshold """ if not stats.get("active"): return False, "inactive" entries = stats.get("entries", 0) old_entries = stats.get("old_entries", stats.get("old-entries", 0)) tombs = stats.get("tombs", 0) total = entries + old_entries if total == 0: return False, "zero-entries" ratio = (tombs + old_entries) / total if ratio <= old_ratio_threshold: return False, f"below-threshold ({ratio:.3f} <= {old_ratio_threshold})" return True, None def hex_key_to_bytes(hex_str): """Decode a PD hex-encoded key string to raw bytes.""" if not hex_str: return b"" try: return bytes.fromhex(hex_str) except ValueError: return b"" def write_checkpoint(path, data): """Atomically write checkpoint JSON to path via a temp file + rename.""" dir_ = os.path.dirname(os.path.abspath(path)) fd, tmp_path = tempfile.mkstemp(dir=dir_, prefix=".compact_ckpt_") try: with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) os.replace(tmp_path, path) except Exception: try: os.unlink(tmp_path) except OSError: pass raise def load_checkpoint(path): """Load and return checkpoint dict, or None if file does not exist.""" try: with open(path, "r", encoding="utf-8") as f: return json.load(f) except FileNotFoundError: return None def build_checkpoint(scan_key_raw, last_region_id, last_start_key, last_end_key, counters): return { "version": CHECKPOINT_VERSION, "timestamp": datetime.now(timezone.utc).isoformat(), "scan_key": scan_key_raw.hex() if scan_key_raw else "", "last_region_id": last_region_id, "last_start_key": last_start_key, "last_end_key": last_end_key, "counters": counters, } def run_compaction(args, pd_addrs, ssl_ctx): use_tls = bool(ssl_ctx) store_cache = {} counters = { "regions_seen": 0, "stats_failures": 0, "leader_store_failures": 0, "inactive": 0, "zero_entries": 0, "below_threshold": 0, "dry_run_would_trigger": 0, "triggered": 0, "trigger_failures": 0, "checkpoint_writes": 0, "checkpoint_loads": 0, } start_key_raw = bytes.fromhex(args.start_key) if args.start_key else b"" end_key_raw = bytes.fromhex(args.end_key) if args.end_key else b"" scan_key_raw = start_key_raw last_region_id = None last_start_key = "" last_end_key = "" if args.checkpoint_file: ckpt = load_checkpoint(args.checkpoint_file) if ckpt: logger.info( "Loaded checkpoint from %s (timestamp=%s, scan_key=%s)", args.checkpoint_file, ckpt.get("timestamp"), ckpt.get("scan_key"), ) scan_key_raw = bytes.fromhex(ckpt["scan_key"]) if ckpt.get("scan_key") else b"" last_region_id = ckpt.get("last_region_id") last_start_key = ckpt.get("last_start_key", "") last_end_key = ckpt.get("last_end_key", "") saved = ckpt.get("counters", {}) counters.update({k: saved.get(k, 0) for k in counters}) counters["checkpoint_loads"] += 1 else: logger.info("No checkpoint found at %s, starting from the beginning", args.checkpoint_file) fatal_pd_error = False interrupted = False last_checkpoint_time = time.monotonic() while True: if end_key_raw and scan_key_raw >= end_key_raw: logger.info("Scan key has reached or passed end key; stopping.") break try: if args.scan_interval > 0: time.sleep(args.scan_interval) regions = scan_regions_batch( pd_addrs, scan_key_raw, end_key_raw, args.region_limit, ssl_ctx, args.timeout, ) except KeyboardInterrupt: logger.info("Interrupted during PD scan; will checkpoint current position.") interrupted = True break except Exception as exc: logger.error("Fatal PD traversal error at scan_key=%s: %s", scan_key_raw.hex(), exc) fatal_pd_error = True break if not regions: logger.info("PD returned no regions; scan complete.") break batch_advanced = False try: for region in regions: region_id = region.get("id") r_start_hex = region.get("start_key", "") r_end_hex = region.get("end_key", "") if region_id is None: logger.warning("Region with missing id, skipping: %s", region) continue counters["regions_seen"] += 1 leader = region.get("leader") if not leader: logger.warning("Region %s has no leader, skipping", region_id) counters["leader_store_failures"] += 1 continue leader_store_id = leader.get("store_id") if not leader_store_id: logger.warning("Region %s leader has no store_id, skipping", region_id) counters["leader_store_failures"] += 1 continue try: base_url = get_store_status_address( pd_addrs, leader_store_id, ssl_ctx, args.timeout, use_tls, store_cache ) except Exception as exc: logger.warning( "Cannot resolve status address for store %s (region %s): %s", leader_store_id, region_id, exc, ) counters["leader_store_failures"] += 1 continue try: stats = get_region_engine_stats(base_url, region_id, ssl_ctx, args.timeout) except Exception as exc: logger.warning("Failed to get engine stats for region %s from %s: %s", region_id, base_url, exc) counters["stats_failures"] += 1 continue eligible, reason = is_eligible(stats, args.old_ratio) if not eligible: if reason == "inactive": counters["inactive"] += 1 elif reason == "zero-entries": counters["zero_entries"] += 1 else: counters["below_threshold"] += 1 logger.debug("Skipping region %s: %s %s", region_id, reason, stats) continue if args.dry_run: counters["dry_run_would_trigger"] += 1 logger.info( "[dry-run] Would compact region %s on store %s (%s) " "(entries=%s, old_entries=%s, tombs=%s)", region_id, leader_store_id, base_url, stats.get("entries"), stats.get("old_entries"), stats.get("tombs"), ) else: compact_url = f"{base_url}/major-compact" try: status, body = http_post( compact_url, ssl_ctx, args.timeout, params={"major_compact": "true", "region_id": region_id}, ) if status == 200: counters["triggered"] += 1 logger.info( "Compaction triggered for region %s on store %s (%s) " "(entries=%s, old_entries=%s, tombs=%s)", region_id, leader_store_id, base_url, stats.get("entries"), stats.get("old_entries"), stats.get("tombs"), ) else: counters["trigger_failures"] += 1 logger.warning( "Compaction POST returned HTTP %s for region %s: %s " "(entries=%s, old_entries=%s, tombs=%s)", status, region_id, body[:200], stats.get("entries"), stats.get("old_entries"), stats.get("tombs"), ) except Exception as exc: counters["trigger_failures"] += 1 logger.warning( "Compaction POST failed for region %s: %s " "(entries=%s, old_entries=%s, tombs=%s)", region_id, exc, stats.get("entries"), stats.get("old_entries"), stats.get("tombs"), ) last_region_id = region_id last_start_key = r_start_hex last_end_key = r_end_hex if args.interval > 0 and not args.dry_run: time.sleep(args.interval) except KeyboardInterrupt: logger.info("Interrupted during region processing; will checkpoint current position.") interrupted = True if interrupted: break next_scan_key = None if regions: last_region = regions[-1] r_end_hex = last_region.get("end_key", "") next_scan_key = hex_key_to_bytes(r_end_hex) if next_scan_key == b"": logger.info("Last region has empty end_key (end of keyspace); stopping.") batch_advanced = True scan_key_raw = b"" break if next_scan_key <= scan_key_raw: logger.error( "Region range not advancing: scan_key=%s, last end_key=%s. " "Possible malformed region data; stopping to prevent infinite loop.", scan_key_raw.hex(), next_scan_key.hex(), ) fatal_pd_error = True break batch_advanced = True scan_key_raw = next_scan_key if batch_advanced and args.checkpoint_file and not args.dry_run: now = time.monotonic() due = (now - last_checkpoint_time) >= args.checkpoint_interval if due: try: ckpt_data = build_checkpoint( scan_key_raw, last_region_id, last_start_key, last_end_key, counters ) write_checkpoint(args.checkpoint_file, ckpt_data) counters["checkpoint_writes"] += 1 last_checkpoint_time = now logger.debug( "Checkpoint written to %s (scan_key=%s)", args.checkpoint_file, scan_key_raw.hex(), ) except Exception as exc: logger.error("Failed to write checkpoint to %s: %s", args.checkpoint_file, exc) if len(regions) < args.region_limit: logger.info("PD returned fewer regions than limit; scan complete.") break if interrupted: logger.info("Run interrupted at scan_key=%s.", scan_key_raw.hex()) logger.info( "Summary: regions_seen=%d, triggered=%d, dry_run_would_trigger=%d, " "trigger_failures=%d, stats_failures=%d, leader_store_failures=%d, " "inactive=%d, zero_entries=%d, below_threshold=%d, " "checkpoint_writes=%d, checkpoint_loads=%d", counters["regions_seen"], counters["triggered"], counters["dry_run_would_trigger"], counters["trigger_failures"], counters["stats_failures"], counters["leader_store_failures"], counters["inactive"], counters["zero_entries"], counters["below_threshold"], counters["checkpoint_writes"], counters["checkpoint_loads"], ) if args.checkpoint_file and interrupted and not args.dry_run: # Write a resume checkpoint on interrupt so the user can restart from # the current position. Skipped during dry-run (read-only mode). try: ckpt_data = build_checkpoint( scan_key_raw, last_region_id, last_start_key, last_end_key, counters ) ckpt_data["interrupted"] = True write_checkpoint(args.checkpoint_file, ckpt_data) logger.info( "Interrupt checkpoint written to %s (scan_key=%s). " "Re-run with the same --checkpoint-file to resume.", args.checkpoint_file, scan_key_raw.hex(), ) except Exception as exc: logger.error("Failed to write interrupt checkpoint: %s", exc) elif args.checkpoint_file and not args.dry_run and not fatal_pd_error: try: ckpt_data = build_checkpoint( scan_key_raw, last_region_id, last_start_key, last_end_key, counters ) ckpt_data["completed"] = True write_checkpoint(args.checkpoint_file, ckpt_data) logger.info( "Compaction run complete. Checkpoint marked complete at %s. " "You may delete this file before the next run to restart from scratch.", args.checkpoint_file, ) except Exception as exc: logger.error("Failed to write final checkpoint: %s", exc) had_errors = ( counters["trigger_failures"] > 0 or counters["stats_failures"] > 0 or counters["leader_store_failures"] > 0 or fatal_pd_error ) if interrupted: return 130 return 1 if had_errors else 0 def parse_args(argv=None): parser = argparse.ArgumentParser( description=( "Trigger major compaction on cloud-storage-engine regions by scanning " "a key range via PD." ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Key format: --start-key and --end-key accept lowercase hex-encoded byte strings. Examples: 41 = byte 0x41 ('A') 4100 = bytes 0x41 0x00 7a = byte 0x7a ('z') Empty start-key means the beginning of the keyspace. Empty end-key means the end of the keyspace. Examples: # Compact all regions from key 0x41 ('A') to key 0x5a ('Z'): python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 \\ --start-key 41 --end-key 5a # Compact the full keyspace: python3 manual_compaction.py --pd http://127.0.0.1:2379 # Dry run: python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 \\ --start-key 41 --end-key 5a --dry-run # Resume from a checkpoint: python3 manual_compaction.py \\ --pd http://127.0.0.1:2379 \\ --start-key 41 --end-key 5a \\ --checkpoint-file /tmp/ckpt.json # TLS with client certs: python3 manual_compaction.py \\ --pd https://127.0.0.1:2379 \\ --ca-cert /var/lib/pd-tls/ca.crt \\ --client-cert /var/lib/pd-tls/tls.crt \\ --client-key /var/lib/pd-tls/tls.key """, ) parser.add_argument( "--pd", dest="pd", default="http://127.0.0.1:2379", help="Comma-separated PD addresses (default: http://127.0.0.1:2379)", ) key_group = parser.add_argument_group("Key range (range-based traversal)") key_group.add_argument( "--start-key", default="", help="Inclusive start key as hex string (e.g. 4100); empty = beginning of keyspace", ) key_group.add_argument( "--end-key", default="", help="Exclusive end key as hex string (e.g. 5a00); empty = end of keyspace", ) key_group.add_argument( "--region-limit", type=int, default=DEFAULT_REGION_LIMIT, help=f"Number of regions to fetch per PD request (default: {DEFAULT_REGION_LIMIT})", ) compact_group = parser.add_argument_group("Compaction eligibility") compact_group.add_argument( "--old-ratio", type=float, default=DEFAULT_OLD_RATIO, help=( f"Trigger compaction when (tombs+old_entries)/(entries+old_entries) > this value " f"(default: {DEFAULT_OLD_RATIO})" ), ) compact_group.add_argument( "--interval", type=float, default=DEFAULT_INTERVAL, help=f"Seconds to sleep between region compaction requests (default: {DEFAULT_INTERVAL})", ) compact_group.add_argument( "--scan-interval", type=float, default=DEFAULT_SCAN_INTERVAL, help=( f"Seconds to sleep between PD scan batches (default: {DEFAULT_SCAN_INTERVAL}). " "Prevents tight-loop CPU pressure on the script and PD when most regions are " "ineligible. Set to 0 to disable." ), ) compact_group.add_argument( "--dry-run", action="store_true", help="Evaluate eligibility and log what would be triggered, but do not POST or modify checkpoints", ) ckpt_group = parser.add_argument_group("Checkpoint") ckpt_group.add_argument( "--checkpoint-file", default="", help=( "Path to a JSON checkpoint file. On startup, if this file exists, " "the scan resumes from the stored key. The file is updated " "periodically (no more often than --checkpoint-interval seconds), " "and also on interrupt and completion. Not written during --dry-run." ), ) ckpt_group.add_argument( "--checkpoint-interval", type=float, default=DEFAULT_CHECKPOINT_INTERVAL, help=( f"Minimum seconds between checkpoint writes (default: {DEFAULT_CHECKPOINT_INTERVAL}). " "Set to 0 to write after every batch." ), ) tls_group = parser.add_argument_group("TLS") tls_group.add_argument("--ca-cert", default="", help="CA certificate file for TLS connections") tls_group.add_argument("--client-cert", default="", help="Client certificate file for mTLS") tls_group.add_argument("--client-key", default="", help="Client private key file for mTLS") misc_group = parser.add_argument_group("Misc") misc_group.add_argument( "--timeout", type=int, default=DEFAULT_TIMEOUT, help=f"HTTP request timeout in seconds (default: {DEFAULT_TIMEOUT})", ) misc_group.add_argument( "--verbose", "-v", action="store_true", help="Enable debug logging", ) return parser.parse_args(argv) def main(argv=None): args = parse_args(argv) if args.verbose: logging.getLogger().setLevel(logging.DEBUG) pd_addrs = [a.strip().rstrip("/") for a in args.pd.split(",") if a.strip()] if not pd_addrs: logger.error("No PD addresses provided") sys.exit(1) ssl_ctx = make_ssl_context(args) if args.dry_run: logger.info("DRY RUN mode: no compaction requests will be issued and checkpoints will not be modified") logger.info( "Starting range-based compaction: pd=%s start_key=%r end_key=%r " "region_limit=%d old_ratio=%.2f dry_run=%s", pd_addrs, args.start_key, args.end_key, args.region_limit, args.old_ratio, args.dry_run, ) rc = run_compaction(args, pd_addrs, ssl_ctx) sys.exit(rc) if __name__ == "__main__": main()