#!/usr/bin/env python3 """ Collect keyspace-level shard stats from TiKV leaders. The script scans PD region metadata for the API v2 transaction key range of a keyspace, resolves each region leader's store status address, fetches /kvengine/ from that leader, and aggregates the per-shard stats into one JSON object. Examples: python3 keyspace_shard_stats.py 10 --pd http://127.0.0.1:2379 # Collect region-level shard stats for one table prefix range: python3 keyspace_shard_stats.py 10 123 --pd http://127.0.0.1:2379 # Collect region-level shard stats for one table index prefix range: python3 keyspace_shard_stats.py 10 --table-id 123 --index-id 5 With TLS/mTLS: python3 keyspace_shard_stats.py 10 \\ --pd https://127.0.0.1:2379 \\ --ca-cert /var/lib/tikv-tls/ca.crt \\ --client-cert /var/lib/tikv-tls/tls.crt \\ --client-key /var/lib/tikv-tls/tls.key """ import argparse import json import logging import ssl import sys import time import urllib.error import urllib.parse import urllib.request PD_REGIONS_KEY_PATH = "/pd/api/v1/regions/key" PD_STORE_PATH = "/pd/api/v1/store" DEFAULT_REGION_LIMIT = 512 DEFAULT_TIMEOUT = 30 DEFAULT_SCAN_INTERVAL = 0.0 TXN_KEY_PREFIX = b"x"[0] MAX_ENCODED_KEYSPACE_ID = 0x00FFFFFE ENC_GROUP_SIZE = 8 ENC_MARKER = 0xFF ENC_PAD = 0x00 I64_MIN = -(1 << 63) I64_MAX = (1 << 63) - 1 I64_SIGN_MARK = 1 << 63 SUM_FIELDS = ( "mem-table-count", "mem-table-size", "mem-table-unpersisted-props-size", "l0-table-count", "blob-table-count", "in-use-blob-size", "total-blob-size", "l0-table-size", "index-size", "filter-size", "entries", "old-entries", "old-data-size", "tombs", "kv-size", "keyspace-prefix-tables", "total-size", "partial-l0s", "shared-blob-tables", "partial-tbls", "txn-file-locks", "columnar-tables", "unconverted-l0-count", ) MAX_FIELDS = ( "max-ts", "base-version", "meta-sequence", "write-sequence", "schema-version", "schema-restore-version", "checked-schema-version", ) BOOL_COUNT_FIELDS = { "active": "num-active-shards", "compacting": "num-compacting-shards", "flushed": "num-initial-flushed-shard", "has-over-bound-data": "num-has-over-bound-data-shards", "has-del-prefixes": "num-has-del-prefixes-shards", "ready-to-destroy-range": "num-ready-destroy-range-shards", "trim-over-bound": "num-trim-over-bound-shards", "manual-major-compaction": "num-manual-major-compaction-shards", "is-sync": "num-sync-shards", "encryption": "num-encrypted-shards", } LEVEL_SUM_FIELDS = ( "num-tables", "data-size", "index-size", "filter-size", "entries", "old-entries", "tombs", "kv-size", "in-use-blob-size", ) COLUMNAR_SUM_FIELDS = ("num-files", "data-size", "data-kv-size") FTS_SUM_FIELDS = ("num-files", "data-size", "pk-total") VECTOR_SUM_FIELDS = ("num-files", "data-size") IA_SUM_FIELDS = ("num-shards", "num-tables", "data-size", "kv-size") logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", datefmt="%Y-%m-%dT%H:%M:%S", stream=sys.stderr, ) 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): if addr.startswith("http://") or addr.startswith("https://"): return addr.rstrip("/") scheme = "https" if use_tls else "http" return f"{scheme}://{addr}".rstrip("/") def _build_opener(ssl_ctx): handler = ( urllib.request.HTTPSHandler(context=ssl_ctx) if ssl_ctx else urllib.request.HTTPHandler() ) return urllib.request.build_opener(handler) def pd_request(pd_addrs, path, ssl_ctx, timeout, params=None): 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) with _build_opener(ssl_ctx).open(req, timeout=timeout) as resp: return json.loads(resp.read()) except Exception as exc: logger.debug("PD request failed for %s: %s", url, exc) last_err = exc raise RuntimeError(f"All PD addresses failed for {path}: {last_err}") def http_get_json(url, ssl_ctx, timeout): try: req = urllib.request.Request(url) with _build_opener(ssl_ctx).open(req, timeout=timeout) as resp: return json.loads(resp.read()) except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise RuntimeError(f"HTTP {exc.code} from {url}: {body}") from exc def encode_txn_keyspace_prefix(keyspace_id): if keyspace_id < 0 or keyspace_id > MAX_ENCODED_KEYSPACE_ID: raise ValueError( f"keyspace_id must be in [0, {MAX_ENCODED_KEYSPACE_ID}] " "for the 3-byte API v2 keyspace prefix" ) encoded = bytearray(keyspace_id.to_bytes(4, byteorder="big")) encoded[0] = TXN_KEY_PREFIX return bytes(encoded) def encode_txn_keyspace_range(keyspace_id): start = encode_txn_keyspace_prefix(keyspace_id) end = encode_txn_keyspace_prefix(keyspace_id + 1) return start, end def encode_i64_comparable(value): if value < I64_MIN or value > I64_MAX: raise ValueError(f"value must be in i64 range: {value}") unsigned = (value & ((1 << 64) - 1)) ^ I64_SIGN_MARK return unsigned.to_bytes(8, byteorder="big") def encode_table_prefix_key(table_id): return b"t" + encode_i64_comparable(table_id) def encode_index_prefix_key(table_id, index_id): return encode_table_prefix_key(table_id) + b"_i" + encode_i64_comparable(index_id) def prefix_next(prefix): if not prefix: raise ValueError("empty prefix has no finite prefix-next key") next_key = bytearray(prefix) for idx in range(len(next_key) - 1, -1, -1): if next_key[idx] != 0xFF: next_key[idx] += 1 del next_key[idx + 1 :] return bytes(next_key) raise ValueError(f"prefix has no finite prefix-next key: {prefix.hex()}") def encode_table_range(keyspace_id, table_id): if table_id == I64_MAX: raise ValueError("table_id must be less than i64::MAX to build a table range") keyspace_prefix = encode_txn_keyspace_prefix(keyspace_id) table_start = encode_table_prefix_key(table_id) table_end = encode_table_prefix_key(table_id + 1) return keyspace_prefix + table_start, keyspace_prefix + table_end def encode_index_range(keyspace_id, table_id, index_id): keyspace_prefix = encode_txn_keyspace_prefix(keyspace_id) index_prefix = encode_index_prefix_key(table_id, index_id) return keyspace_prefix + index_prefix, keyspace_prefix + prefix_next(index_prefix) def encode_pd_region_key(raw_key): encoded = bytearray() for offset in range(0, len(raw_key) + 1, ENC_GROUP_SIZE): group = raw_key[offset : offset + ENC_GROUP_SIZE] pad_size = ENC_GROUP_SIZE - len(group) encoded.extend(group) encoded.extend(bytes([ENC_PAD]) * pad_size) encoded.append(ENC_MARKER - pad_size) if pad_size: break return bytes(encoded) def encode_txn_keyspace_pd_range(keyspace_id): start, end = encode_txn_keyspace_range(keyspace_id) return encode_pd_region_key(start), encode_pd_region_key(end) def build_scope_keys(keyspace_id, table_id=None, index_id=None): if index_id is not None: if table_id is None: raise ValueError("index_id requires table_id") start_key, end_key = encode_index_range(keyspace_id, table_id, index_id) scope = "index" elif table_id is not None: start_key, end_key = encode_table_range(keyspace_id, table_id) scope = "table" else: start_key, end_key = encode_txn_keyspace_range(keyspace_id) scope = "keyspace" return scope, start_key, end_key, encode_pd_region_key(start_key), encode_pd_region_key(end_key) def hex_key_to_bytes(hex_str): if not hex_str: return b"" return bytes.fromhex(hex_str) def scan_regions_batch(pd_addrs, start_key_raw, end_key_raw, limit, ssl_ctx, timeout): params = { "key": start_key_raw.hex() if start_key_raw else "", "end_key": end_key_raw.hex() if end_key_raw else "", "limit": str(limit), "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): if store_id in store_cache: return store_cache[store_id] data = pd_request(pd_addrs, f"{PD_STORE_PATH}/{store_id}", ssl_ctx, timeout) 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") base_url = normalize_status_address(raw_addr, use_tls) store_cache[store_id] = base_url return base_url def get_region_shard_stats(base_url, region_id, ssl_ctx, timeout): return http_get_json(f"{base_url}/kvengine/{region_id}", ssl_ctx, timeout) def _number(value): if isinstance(value, bool) or value is None: return 0 if isinstance(value, (int, float)): return value return 0 def _sum_into(dst, src, fields): for field in fields: dst[field] = dst.get(field, 0) + _number(src.get(field)) def _max_into(dst, src, src_field, dst_field=None): dst_field = dst_field or src_field dst[dst_field] = max(dst.get(dst_field, 0), _number(src.get(src_field))) def _ensure_list_len(values, length, factory): while len(values) < length: values.append(factory()) def _merge_numeric_list(dst, field, src): if not isinstance(src, list): return values = dst.setdefault(field, []) _ensure_list_len(values, len(src), lambda: 0) for idx, value in enumerate(src): values[idx] += _number(value) def _merge_level_stats(dst, src): if "level" not in dst: dst["level"] = src.get("level", 0) _sum_into(dst, src, LEVEL_SUM_FIELDS) _max_into(dst, src, "max-ts") def _merge_cfs(dst, src_cfs): if not isinstance(src_cfs, list): return dst_cfs = dst.setdefault("cfs", []) cf_level_files = dst.setdefault("cf-level-files", []) cf_level_sizes = dst.setdefault("cf-level-sizes", []) _ensure_list_len(dst_cfs, len(src_cfs), lambda: {"levels": []}) _ensure_list_len(cf_level_files, len(src_cfs), list) _ensure_list_len(cf_level_sizes, len(src_cfs), list) for cf_idx, src_cf in enumerate(src_cfs): src_levels = src_cf.get("levels", []) if isinstance(src_cf, dict) else [] dst_levels = dst_cfs[cf_idx].setdefault("levels", []) _ensure_list_len(dst_levels, len(src_levels), dict) _ensure_list_len(cf_level_files[cf_idx], len(src_levels), lambda: 0) _ensure_list_len(cf_level_sizes[cf_idx], len(src_levels), lambda: 0) for level_idx, src_level in enumerate(src_levels): if not isinstance(src_level, dict): continue _merge_level_stats(dst_levels[level_idx], src_level) cf_level_files[cf_idx][level_idx] += _number(src_level.get("num-tables")) cf_level_sizes[cf_idx][level_idx] += _number(src_level.get("data-size")) def _merge_indexed_dict_list(dst, field, src, sum_fields): if not isinstance(src, list): return values = dst.setdefault(field, []) _ensure_list_len(values, len(src), dict) for idx, src_item in enumerate(src): if isinstance(src_item, dict): _sum_into(values[idx], src_item, sum_fields) def _merge_dict(dst, field, src, sum_fields): if not isinstance(src, dict): return values = dst.setdefault(field, {}) _sum_into(values, src, sum_fields) def new_aggregate_stats(): stats = { "num-shards": 0, "num-pending-compaction-shards": 0, "ready-destroy-range-shards": [], "max-compaction-score": 0.0, "storage-class-spec-counts": {}, } for field in SUM_FIELDS: stats[field] = 0 for field in MAX_FIELDS: stats[f"max-{field}"] = 0 for count_field in BOOL_COUNT_FIELDS.values(): stats[count_field] = 0 return stats def aggregate_shard_stats(aggregate, shard_stats): aggregate["num-shards"] += 1 for field, count_field in BOOL_COUNT_FIELDS.items(): if shard_stats.get(field): aggregate[count_field] += 1 if _number(shard_stats.get("compaction-score")) > 1.0 and not shard_stats.get("compacting"): aggregate["num-pending-compaction-shards"] += 1 if shard_stats.get("ready-to-destroy-range"): aggregate["ready-destroy-range-shards"].append( {"id": shard_stats.get("id", 0), "ver": shard_stats.get("ver", 0)} ) if _number(shard_stats.get("inner-key-off")) > 0: aggregate["num-inner-key-shards"] = aggregate.get("num-inner-key-shards", 0) + 1 _sum_into(aggregate, shard_stats, SUM_FIELDS) for field in MAX_FIELDS: _max_into(aggregate, shard_stats, field, f"max-{field}") _max_into(aggregate, shard_stats, "compaction-score", "max-compaction-score") storage_class_spec = shard_stats.get("storage-class-spec", "") if storage_class_spec: counts = aggregate["storage-class-spec-counts"] counts[storage_class_spec] = counts.get(storage_class_spec, 0) + 1 _merge_numeric_list(aggregate, "l0-cf-table-size", shard_stats.get("l0-cf-table-size")) _merge_cfs(aggregate, shard_stats.get("cfs")) _merge_indexed_dict_list( aggregate, "columnar-levels", shard_stats.get("columnar-levels"), COLUMNAR_SUM_FIELDS, ) _merge_indexed_dict_list( aggregate, "fts-levels", shard_stats.get("fts-levels"), FTS_SUM_FIELDS, ) _merge_dict(aggregate, "vector-indexes", shard_stats.get("vector-indexes"), VECTOR_SUM_FIELDS) _merge_dict(aggregate, "ia", shard_stats.get("ia"), IA_SUM_FIELDS) total_data_size = shard_stats.get("l0-table-size", 0) for cf in shard_stats.get("cfs", []) or []: for level in cf.get("levels", []) if isinstance(cf, dict) else []: total_data_size += _number(level.get("data-size")) for level in shard_stats.get("columnar-levels", []) or []: total_data_size += _number(level.get("data-size")) for level in shard_stats.get("fts-levels", []) or []: total_data_size += _number(level.get("data-size")) total_data_size += _number((shard_stats.get("vector-indexes") or {}).get("data-size")) aggregate["total-data-size"] = aggregate.get("total-data-size", 0) + total_data_size def _new_result(keyspace_id, table_id, index_id, scope, start_key, end_key, include_regions): result = { "scope": scope, "keyspace-id": keyspace_id, "start-key": start_key.hex(), "end-key": end_key.hex(), "regions": { "seen": 0, "summed": 0, "range-boundary-overlaps": 0, "malformed": 0, "without-leader": 0, "leader-store-failures": 0, "stats-failures": 0, "mismatched-stats": 0, }, "stores": {}, "stats": new_aggregate_stats(), "complete": True, "range-exact-by-region-boundaries": True, } if table_id is not None: result["table-id"] = table_id if index_id is not None: result["index-id"] = index_id if include_regions: result["region-details"] = [] return result def _record_store(result, store_id, status_address): store_key = str(store_id) stores = result["stores"] if store_key not in stores: stores[store_key] = {"status-address": status_address, "regions": 0} stores[store_key]["regions"] += 1 def _validate_shard_stats(result, region_id, keyspace_id, shard_stats): stats_id = shard_stats.get("id") if stats_id != region_id: logger.warning("Shard stats id mismatch: region=%s stats_id=%s", region_id, stats_id) result["regions"]["mismatched-stats"] += 1 return False stats_keyspace = shard_stats.get("keyspace") if stats_keyspace is not None and stats_keyspace != keyspace_id: logger.warning( "Shard stats keyspace mismatch: region=%s expected=%s actual=%s", region_id, keyspace_id, stats_keyspace, ) result["regions"]["mismatched-stats"] += 1 return False return True def _region_extends_outside_scope(region, scope_start_key, scope_end_key): region_start = hex_key_to_bytes(region.get("start_key", "")) region_end_hex = region.get("end_key", "") region_end = hex_key_to_bytes(region_end_hex) if region_end_hex else None return region_start < scope_start_key or region_end is None or region_end > scope_end_key def collect_keyspace_stats(args, pd_addrs, ssl_ctx): use_tls = bool(ssl_ctx) store_cache = {} scope, start_key, end_key, pd_start_key, pd_end_key = build_scope_keys( args.keyspace_id, args.table_id, args.index_id ) result = _new_result( args.keyspace_id, args.table_id, args.index_id, scope, start_key, end_key, args.include_regions, ) result["pd-start-key"] = pd_start_key.hex() result["pd-end-key"] = pd_end_key.hex() scan_key = pd_start_key while True: if scan_key >= pd_end_key: break if args.scan_interval > 0: time.sleep(args.scan_interval) try: regions = scan_regions_batch( pd_addrs, scan_key, pd_end_key, args.region_limit, ssl_ctx, args.timeout, ) except Exception as exc: logger.error("PD traversal failed at scan_key=%s: %s", scan_key.hex(), exc) result["complete"] = False result["fatal-error"] = str(exc) return result if not regions: break for region in regions: region_id = region.get("id") if region_id is None: logger.warning("Region with missing id, skipping: %s", region) result["regions"]["malformed"] += 1 continue result["regions"]["seen"] += 1 range_boundary_overlap = _region_extends_outside_scope(region, pd_start_key, pd_end_key) if range_boundary_overlap: result["regions"]["range-boundary-overlaps"] += 1 result["range-exact-by-region-boundaries"] = False leader = region.get("leader") if not leader: logger.warning("Region %s has no leader, skipping", region_id) result["regions"]["without-leader"] += 1 continue leader_store_id = leader.get("store_id") if leader_store_id is None: logger.warning("Region %s leader has no store_id, skipping", region_id) result["regions"]["without-leader"] += 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, ) result["regions"]["leader-store-failures"] += 1 continue _record_store(result, leader_store_id, base_url) try: shard_stats = get_region_shard_stats(base_url, region_id, ssl_ctx, args.timeout) except Exception as exc: logger.warning( "Failed to get shard stats for region %s from %s: %s", region_id, base_url, exc, ) result["regions"]["stats-failures"] += 1 continue if not _validate_shard_stats(result, region_id, args.keyspace_id, shard_stats): continue aggregate_shard_stats(result["stats"], shard_stats) result["regions"]["summed"] += 1 if args.include_regions: result["region-details"].append( { "region-id": region_id, "leader-store-id": leader_store_id, "store-status-address": base_url, "start-key": region.get("start_key", ""), "end-key": region.get("end_key", ""), "range-boundary-overlap": range_boundary_overlap, } ) last_region = regions[-1] last_end_key = last_region.get("end_key", "") if not last_end_key: break next_scan_key = hex_key_to_bytes(last_end_key) if next_scan_key <= scan_key: result["complete"] = False result["fatal-error"] = ( f"Region range did not advance: scan_key={scan_key.hex()} " f"last_end_key={last_end_key}" ) logger.error(result["fatal-error"]) return result scan_key = next_scan_key if len(regions) < args.region_limit: break failures = sum( result["regions"][field] for field in ( "malformed", "without-leader", "leader-store-failures", "stats-failures", "mismatched-stats", ) ) if failures: result["complete"] = False return result def parse_args(argv=None): parser = argparse.ArgumentParser( description="Collect aggregated shard stats for one API v2 keyspace.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Output: The script writes one JSON object to stdout. Warnings and progress logs go to stderr. If any region cannot be collected, the JSON is still emitted with complete=false and the process exits non-zero unless --allow-partial is set. """, ) parser.add_argument("keyspace_id", type=int, help="Keyspace ID to collect") parser.add_argument( "table_id", nargs="?", type=int, help="Optional table ID. When set, collect the table prefix range.", ) parser.add_argument( "--pd", default="http://127.0.0.1:2379", help="Comma-separated PD addresses (default: http://127.0.0.1:2379)", ) parser.add_argument( "--table-id", dest="table_id_option", type=int, help="Optional table ID. Equivalent to the positional table_id argument.", ) parser.add_argument( "--index-id", type=int, help="Optional index ID. Requires table_id and collects the index prefix range.", ) parser.add_argument( "--region-limit", type=int, default=DEFAULT_REGION_LIMIT, help=f"Number of regions to fetch per PD request (default: {DEFAULT_REGION_LIMIT})", ) parser.add_argument( "--scan-interval", type=float, default=DEFAULT_SCAN_INTERVAL, help=f"Seconds to sleep between PD scan batches (default: {DEFAULT_SCAN_INTERVAL})", ) parser.add_argument( "--timeout", type=int, default=DEFAULT_TIMEOUT, help=f"HTTP request timeout in seconds (default: {DEFAULT_TIMEOUT})", ) parser.add_argument("--ca-cert", default="", help="CA certificate file for TLS connections") parser.add_argument("--client-cert", default="", help="Client certificate file for mTLS") parser.add_argument("--client-key", default="", help="Client private key file for mTLS") parser.add_argument( "--include-regions", action="store_true", help="Include collected region ids and leader stores in the output", ) parser.add_argument( "--allow-partial", action="store_true", help="Exit 0 even when some regions could not be collected", ) parser.add_argument("--verbose", "-v", action="store_true", help="Enable debug logging") args = parser.parse_args(argv) if args.table_id is not None and args.table_id_option is not None: if args.table_id != args.table_id_option: parser.error("positional table_id and --table-id must be the same when both are set") elif args.table_id is None: args.table_id = args.table_id_option if args.index_id is not None and args.table_id is None: parser.error("--index-id requires table_id or --table-id") return args def main(argv=None): args = parse_args(argv) if args.verbose: logging.getLogger().setLevel(logging.DEBUG) try: build_scope_keys(args.keyspace_id, args.table_id, args.index_id) except ValueError as exc: logger.error("%s", exc) return 2 if args.region_limit <= 0: logger.error("--region-limit must be greater than 0") return 2 pd_addrs = [addr.strip().rstrip("/") for addr in args.pd.split(",") if addr.strip()] if not pd_addrs: logger.error("No PD addresses provided") return 2 ssl_ctx = make_ssl_context(args) result = collect_keyspace_stats(args, pd_addrs, ssl_ctx) json.dump(result, sys.stdout, indent=2) sys.stdout.write("\n") if result.get("complete") or args.allow_partial: return 0 return 1 if __name__ == "__main__": sys.exit(main())