diff --git a/scripts/bpf/trace.py b/scripts/bpf/trace.py index f5ad7e164..450016134 100755 --- a/scripts/bpf/trace.py +++ b/scripts/bpf/trace.py @@ -4,7 +4,129 @@ from argparse import ArgumentParser from dataclasses import dataclass from typing import Dict, List, TypeVar import json +import os +import re +import subprocess import sys +import tempfile + + +@dataclass +class DTraceArgument: + """Describes a DTrace probe (usdt) argument""" + name: str + pos: int + type: type + + +@dataclass +class DTraceProbe: + """Describes a DTrace probe (usdt) point""" + name: str + args: Dict[str, DTraceArgument] + + def __init__(self, name, args): + self.name = name + self.args = {a.name: a for a in args} + + +@dataclass +class DTraceEntry: + """Describes a single DTrace probe invocation""" + name: str + args: Dict[str, TypeVar('ArgumentType', str, int)] + + def __init__(self, probe, args): + valmap = {int: lambda x: int(x, 16), + str: lambda x: x.strip().strip("'")} + self.name = probe.name + self.args = {} + for name, value in args.items(): + arg = probe.args.get(name) + if arg is None: + raise ValueError(f'Unexpected argument: {name}') + self.args[name] = valmap[arg.type](value) + + +class DTrace: + """Generates bpftrace script based on the supplied probe points, parses its + output and stores is as a list of DTraceEntry sorted by their tsc. + """ + def __init__(self, probes, file=None): + self._avail_probes = self._list_probes() + self._probes = {p.name: p for p in probes} + self.entries = self._parse(file) if file is not None else [] + # Sanitize the probe definitions + for probe in probes: + if probe.name not in self._avail_probes: + raise ValueError(f'Couldn\'t find probe: "{probe.name}"') + for arg in probe.args.values(): + if arg.pos >= self._avail_probes[probe.name]: + raise ValueError('Invalid probe argument position') + if arg.type not in (int, str): + raise ValueError('Invalid argument type') + + def _parse(self, file): + regex = re.compile(r'(\w+): (.*)') + entries = [] + + for line in file.readlines(): + match = regex.match(line) + if match is None: + continue + name, args = match.groups() + probe = self._probes.get(name) + # Skip the line if we don't recognize the probe name + if probe is None: + continue + entries.append(DTraceEntry(probe, args=dict(a.strip().split('=') + for a in args.split(',')))) + entries.sort(key=lambda e: e.args['tsc']) + return entries + + def _list_probes(self): + files = subprocess.check_output(['git', 'ls-files', '*.[ch]', + ':!:include/spdk_internal/usdt.h']) + files = filter(lambda f: len(f) > 0, str(files, 'ascii').split('\n')) + regex = re.compile(r'SPDK_DTRACE_PROBE([0-9]*)\((\w+)') + probes = {} + + for fname in files: + with open(fname, 'r') as file: + for match in regex.finditer(file.read()): + nargs, name = match.group(1), match.group(2) + nargs = int(nargs) if len(nargs) > 0 else 0 + # Add one to accommodate for the tsc being the first arg + probes[name] = nargs + 1 + return probes + + def _gen_usdt(self, probe): + usdt = (f'usdt:__EXE__:{probe.name} {{' + + f'printf("{probe.name}: ') + args = probe.args + if len(args) > 0: + argtype = {int: '0x%lx', str: '\'%s\''} + argcast = {int: lambda x: x, str: lambda x: f'str({x})'} + argstr = [f'{a.name}={argtype[a.type]}' for a in args.values()] + argval = [f'{argcast[a.type](f"arg{a.pos}")}' for a in args.values()] + usdt += ', '.join(argstr) + '\\n", ' + ', '.join(argval) + else: + usdt += '\\n"' + usdt += ');}' + return usdt + + def generate(self): + return '\n'.join([self._gen_usdt(p) for p in self._probes.values()]) + + def record(self, pid): + with tempfile.NamedTemporaryFile(mode='w+') as script: + script.write(self.generate()) + script.flush() + try: + subprocess.run([f'{os.path.dirname(__file__)}/../bpftrace.sh', + f'{pid}', f'{script.name}']) + except KeyboardInterrupt: + pass @dataclass @@ -101,14 +223,42 @@ class Trace: print(' '.join([*filter(lambda f: f is not None, fields)]).rstrip()) +def build_dtrace(): + return DTrace([ + DTraceProbe( + name='nvmf_poll_group_add_qpair', + args=[DTraceArgument(name='tsc', pos=0, type=int), + DTraceArgument(name='qpair', pos=1, type=int), + DTraceArgument(name='thread', pos=2, type=int)]), + DTraceProbe( + name='nvmf_poll_group_remove_qpair', + args=[DTraceArgument(name='tsc', pos=0, type=int), + DTraceArgument(name='qpair', pos=1, type=int), + DTraceArgument(name='thread', pos=2, type=int)]), + DTraceProbe( + name='nvmf_ctrlr_add_qpair', + args=[DTraceArgument(name='tsc', pos=0, type=int), + DTraceArgument(name='qpair', pos=1, type=int), + DTraceArgument(name='qid', pos=2, type=int), + DTraceArgument(name='subnqn', pos=3, type=str), + DTraceArgument(name='hostnqn', pos=4, type=str)])]) + + def main(argv): parser = ArgumentParser(description='SPDK trace annotation script') parser.add_argument('-i', '--input', help='JSON-formatted trace file produced by spdk_trace app') + parser.add_argument('-g', '--generate', help='Generate bpftrace script', action='store_true') + parser.add_argument('-r', '--record', help='Record BPF traces on PID', metavar='PID', type=int) args = parser.parse_args(argv) - file = open(args.input, 'r') if args.input is not None else sys.stdin - Trace(file).print() + if args.generate: + print(build_dtrace().generate()) + elif args.record: + build_dtrace().record(args.record) + else: + file = open(args.input, 'r') if args.input is not None else sys.stdin + Trace(file).print() if __name__ == '__main__':