1#!/usr/bin/python
2#
3# dbstat        Display a histogram of MySQL and PostgreSQL query latencies.
4#
5# USAGE: dbstat [-v] [-p PID [PID ...]] [-m THRESHOLD] [-u]
6#               [-i INTERVAL] {mysql,postgres}
7#
8# This tool uses USDT probes, which means it needs MySQL and PostgreSQL built
9# with USDT (DTrace) support.
10#
11# Copyright 2017, Sasha Goldshtein
12# Licensed under the Apache License, Version 2.0
13#
14# 15-Feb-2017   Sasha Goldshtein   Created this.
15
16from bcc import BPF, USDT
17import argparse
18import subprocess
19from time import sleep, strftime
20
21examples = """
22    dbstat postgres     # display a histogram of PostgreSQL query latencies
23    dbstat mysql -v     # display MySQL latencies and print the BPF program
24    dbstat mysql -u     # display query latencies in microseconds (default: ms)
25    dbstat mysql -m 5   # trace only queries slower than 5ms
26    dbstat mysql -p 408 # trace queries in a specific process
27"""
28parser = argparse.ArgumentParser(
29    description="",
30    formatter_class=argparse.RawDescriptionHelpFormatter,
31    epilog=examples)
32parser.add_argument("-v", "--verbose", action="store_true",
33    help="print the BPF program")
34parser.add_argument("db", choices=["mysql", "postgres"],
35    help="the database engine to use")
36parser.add_argument("-p", "--pid", type=int, nargs='*',
37    dest="pids", metavar="PID", help="the pid(s) to trace")
38parser.add_argument("-m", "--threshold", type=int, default=0,
39    help="trace queries slower than this threshold (ms)")
40parser.add_argument("-u", "--microseconds", action="store_true",
41    help="display query latencies in microseconds (default: milliseconds)")
42parser.add_argument("-i", "--interval", type=int, default=99999999999,
43    help="print summary at this interval (seconds)")
44args = parser.parse_args()
45
46if not args.pids or len(args.pids) == 0:
47    if args.db == "mysql":
48        args.pids = map(int, subprocess.check_output(
49                                        "pidof mysqld".split()).split())
50    elif args.db == "postgres":
51        args.pids = map(int, subprocess.check_output(
52                                        "pidof postgres".split()).split())
53
54program = """
55#include <uapi/linux/ptrace.h>
56
57BPF_HASH(temp, u64, u64);
58BPF_HISTOGRAM(latency);
59
60int probe_start(struct pt_regs *ctx) {
61    u64 timestamp = bpf_ktime_get_ns();
62    u64 pid = bpf_get_current_pid_tgid();
63    temp.update(&pid, &timestamp);
64    return 0;
65}
66
67int probe_end(struct pt_regs *ctx) {
68    u64 *timestampp;
69    u64 pid = bpf_get_current_pid_tgid();
70    timestampp = temp.lookup(&pid);
71    if (!timestampp)
72        return 0;
73
74    u64 delta = bpf_ktime_get_ns() - *timestampp;
75    FILTER
76    delta /= SCALE;
77    latency.increment(bpf_log2l(delta));
78    temp.delete(&pid);
79    return 0;
80}
81"""
82program = program.replace("SCALE", str(1000 if args.microseconds else 1000000))
83program = program.replace("FILTER", "" if args.threshold == 0 else
84        "if (delta / 1000000 < %d) { return 0; }" % args.threshold)
85
86usdts = map(lambda pid: USDT(pid=pid), args.pids)
87for usdt in usdts:
88    usdt.enable_probe("query__start", "probe_start")
89    usdt.enable_probe("query__done", "probe_end")
90
91if args.verbose:
92    print('\n'.join(map(lambda u: u.get_text(), usdts)))
93    print(program)
94
95bpf = BPF(text=program, usdt_contexts=usdts)
96
97print("Tracing database queries for pids %s slower than %d ms..." %
98      (', '.join(map(str, args.pids)), args.threshold))
99
100latencies = bpf["latency"]
101
102def print_hist():
103    print("[%s]" % strftime("%H:%M:%S"))
104    latencies.print_log2_hist("query latency (%s)" %
105                              ("us" if args.microseconds else "ms"))
106    print("")
107    latencies.clear()
108
109while True:
110    try:
111        sleep(args.interval)
112        print_hist()
113    except KeyboardInterrupt:
114        print_hist()
115        break
116