schedtime

#!/usr/bin/env python3
#
# Copyright 2020 Matt Fleming
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# schedtime - Display a task run time statistics using scheduler hooks.
#
# Unlike bash's time built-in schedtime includes off-cpu time in its results.
#

from __future__ import print_function
from bcc import BPF
import argparse
import os
import signal
import sys
import subprocess
import time

source="""
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

struct data_item {
    u32 pid;
    u32 tgid;
    int kernel_stack_id;
};

struct pid_item {
    u32 parent;
    u32 child;
};

struct task_lifetime {
    u32 pid;
    u64 delta;
    u32 dead;
    char comm[TASK_COMM_LEN];
};

BPF_HASH(events, struct data_item);
BPF_HASH(start, u32);
BPF_STACK_TRACE(stack_traces, 40960);
BPF_HASH(pids, struct pid_item);
BPF_HASH(task_lifetime, u32, struct task_lifetime);

int wake_up_probe(struct pt_regs *ctx, struct task_struct *p)
{
    struct pid_item pid = {};
    struct task_lifetime tl = {};
    u64 one = 1;
    u64 now;
    u32 child = p->pid;

    // Don't trace kernel threads
    if (p->flags & PF_KTHREAD)
        return 0;

    pid.parent = bpf_get_current_pid_tgid();
    pid.child = child;

    pids.update(&pid, &one);

    now = bpf_ktime_get_ns();
    tl.pid = child;
    tl.delta = now;
    task_lifetime.update(&child, &tl);

    return 0;
}

int sched_in(struct pt_regs *ctx, struct task_struct *prev) {
    struct data_item data = {};
    u32 pid = prev->pid;
    u32 ppid;
    u32 tgid;
    u64 ts, *tsp;
    u64 delta;

    // We may never see 'prev' again if it's a dying task. Update its stats.
    ts = bpf_ktime_get_ns();
    start.update(&pid, &ts);

    // get the current thread's start time
    pid = bpf_get_current_pid_tgid();
    tgid = bpf_get_current_pid_tgid() >> 32;
    tsp = start.lookup(&pid);
    if (tsp == 0)
        return 0;

    // calculate current thread's delta time
    delta = bpf_ktime_get_ns() - *tsp;
    start.delete(&pid);
    delta = delta / 1000;

    data.pid = pid;
    data.tgid = tgid;
    data.kernel_stack_id = stack_traces.get_stackid(ctx, 0);

    events.increment(data, delta);

    return 0;
}

int dead_task(struct pt_regs *ctx, struct task_struct *p)
{
    u64 *begin, delta;
    u64 now = bpf_ktime_get_ns();
    u32 pid = p->pid;
    struct task_lifetime *tl;

    tl = task_lifetime.lookup(&pid);
    if (!tl)
        return 0;

    delta = now - tl->delta;
    if ((s64)delta <= 0)
        delta = 1000;
    tl->delta = delta / 1000;
    tl->dead = 1;
    bpf_get_current_comm(&tl->comm, sizeof(tl->comm));
    task_lifetime.update(&pid, tl);
    return 0;
}
"""

fh = sys.stdout
def eprint(*args, **kwargs):
    print(*args, file=fh, **kwargs)

args = None

def time_unit(us_val):
    return us_val / 1000.0 if args.milliseconds else us_val / 1000.0 / 1000.0

def pct(x, total):
    return 0.0 if x == 0 else x / total * 100.0

# Breadth-first search
#
# Find all vertices reachable from source vertex s.
def BFS(s, edges):
    # Build adjacency lists
    adj = {s: []}
    for k,v in edges:
        if k.parent in adj:
            adj[k.parent].append(k.child)
        else:
            adj[k.parent] = [k.child]

        if k.child not in adj:
            adj[k.child] = []

    vertices = set()
    for k, v in edges:
        vertices.add(k.parent)
        vertices.add(k.child)

    WHITE = 0
    GREY = 1
    BLACK = 2

    colour = {}
    parent = {}

    for u in vertices:
        colour[u] = WHITE
        parent[u] = None

    colour[s] = GREY
    q = [s]
    while q:
        u = q.pop(0)    # FIFO queue
        for v in adj[u]:
            if colour[v] == WHITE:
                colour[v] = GREY
                parent[v] = u
                q.append(v)
        colour[u] = BLACK

    return [u for u in colour if colour[u] == BLACK]

class Task(object):
    """A Linux task (thread)"""
    def __init__(self, pid):
        self.lifetime = 0
        self.pid = pid
        self.comm = ""
        self.times = {
                "disk io": 0,
                "voluntary wait": 0,
                "involuntary preempt": 0,
                "userspace locking": 0,
                "kernel locking": 0,
                "network io": 0,
                "page faults": 0
        }

    def total_time(self):
        """Return the total time this task ran."""
        return self.lifetime - sum([self.times[g] for g in self.times])

    def __str__(self):
        unit = "ms" if args.milliseconds else "s"
        oncpu = { "abs": self.total_time(), "pct": pct(self.total_time(), self.lifetime)}

        return "  [{27:s}-{0:d}] lifetime: {1:.2f}{2:s}\n" \
               "\n" \
               "                   on-cpu: {3: >6.2f}{4:s} ({5:>4.1f}%)\n" \
               "\n" \
               "        userspace locking: {6: >6.2f}{7:s} ({8:>4.1f}%)\n" \
               "           userspace wait: {9: >6.2f}{10:s} ({11:>4.1f}%)\n"\
               "\n" \
               "           kernel locking: {12: >6.2f}{13:s} ({14:>4.1f}%)\n" \
               "              kernel wait: {15: >6.2f}{16:s} ({17:>4.1f}%)\n" \
               "\n" \
               "              network i/o: {18: >6.2f}{19:s} ({20:>4.1f}%)\n" \
               "                 disk i/o: {21: >6.2f}{22:s} ({23:>4.1f}%)\n" \
               "\n" \
               "              page faults: {24: >6.2f}{25:s} ({26:>4.1f}%)\n" \
               .format(
                self.pid, time_unit(self.lifetime), unit, time_unit(oncpu["abs"]), unit, oncpu["pct"],
                time_unit(self.times["userspace locking"]), unit, self.pct("userspace locking"),
                time_unit(self.times["voluntary wait"]), unit, self.pct("voluntary wait"),
                time_unit(self.times["kernel locking"]), unit, self.pct("kernel locking"),
                time_unit(self.times["involuntary preempt"]), unit, self.pct("involuntary preempt"),
                time_unit(self.times["network io"]), unit, self.pct("network io"),
                time_unit(self.times["disk io"]), unit, self.pct("disk io"),
                time_unit(self.times["page faults"]), unit, self.pct("page faults"),
                self.comm,
                )

    def pct(self, group):
        return pct(self.times[group], self.lifetime)


def sort_task_data(tasks, args):
    if not args.sort or args.sort == "lifetime":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.lifetime, reverse=True)

    if args.sort == "command":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.comm, reverse=False)
    if args.sort == "tid":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.pid, reverse=False)
    if args.sort == "cpu":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.total_time(), reverse=True)
    if args.sort == "disk":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["disk io"], reverse=True)
    if args.sort == "faults":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["page faults"], reverse=True)
    if args.sort == "net":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["network io"], reverse=True)
    if args.sort == "klock":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["kernel locking"], reverse=True)
    if args.sort == "kwait":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["involuntary preempt"], reverse=True)
    if args.sort == "ulock":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["userspace locking"], reverse=True)
    if args.sort == "uwait":
        return sorted([tasks[t] for t in tasks],
                key=lambda t: t.times["voluntary wait"], reverse=True)

    print("Invalid sort key")
    return []

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
            description="Summarise scheduler time statistics for tasks",
            formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-v", "--verbose", action="store_true",
            help="display verbose time statistics for each task individually")
    parser.add_argument("-m", "--milliseconds", action="store_true",
            help="display time in milliseconds")
    parser.add_argument("-o", "--output",
            help="write output to file instead of stdout")
    parser.add_argument("-s", "--sort",
            help="sort the output of -v by one of: \n\n"
                 "    lifetime - task lifetime (default),\n"
                 "    command - the task command name,\n"
                 "    tid - the task thread id,\n"
                 "    cpu - on-cpu time,\n"
                 "    disk - disk i/o time,\n"
                 "    faults - page fault time,\n"
                 "    net - network i/o time,\n"
                 "    klock - kernel locking time,\n"
                 "    kwait - kernel wait time,\n"
                 "    ulock - userspace locking time,\n"
                 "    uwait - userspace wait time\n\n"
                 "    e.g. '-s cpu' sorts output by on-cpu time")
    parser.add_argument("command",  help="comamand to execute")
    parser.add_argument("args", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    if args.sort and not args.verbose:
        print("Error: -s requires -v")
        sys.exit(1)

    if args.output:
        fh = open(args.output, 'w')

    b = BPF(text=source)

    b.attach_kprobe(event_re=r'^finish_task_switch$|^finish_task_switch\.isra\.\d$', fn_name="sched_in")
    b.attach_kprobe(event="wake_up_new_task", fn_name="wake_up_probe")
    b.attach_kprobe(event="exit_thread", fn_name="dead_task")

    proc = subprocess.Popen([args.command] + args.args)
    try:
        stderr, stdout = proc.communicate()
    except KeyboardInterrupt:
        pass
    finally:
        proc.wait()

    tasks = {}
    p = b.get_table("pids")
    for p in BFS(proc.pid, p.items()):
        tasks[p] = Task(p)
    
    # Unfortunately some tasks may not be gone through do_exit() by the
    # time we try to read their lifetimes.
    wall_time = 0.0
    lifetimes = b.get_table("task_lifetime")
    ltimes = []
    for k, v in lifetimes.items():
        if k.value in tasks and v.dead:
            t = tasks[k.value]
            t.lifetime = v.delta
            t.comm = v.comm.decode()
            ltimes.append(k.value)
        if k.value == proc.pid:
            wall_time = v.delta

    # XXX This shouldn't be needed. How do we even end up with tasks in
    # the BFS tree if we haven't recorded their lifetime?
    bad_tasks = [t for t in tasks if tasks[t].lifetime == 0]
    for t in bad_tasks:
        del tasks[t]

    events = b.get_table("events")
    stack_traces = b.get_table("stack_traces")
    for k, v in sorted(events.items(), key=lambda counts:
            counts[1].value):

        kernel_stack = [] if k.kernel_stack_id < 0 else \
            stack_traces.walk(k.kernel_stack_id)

        pid = k.pid
        if pid not in tasks:
            continue

        # The call stack should look like this:
        # finish_task_switch()
        # schedule()
        # [exit_to_usermode_loop | io_schedule | do_wait]
        kstack = [b.ksym(addr).decode("utf-8") for addr in kernel_stack]

        if len(kstack) < 3:
            continue

        # Some stacks don't have enough information to analyse.
        # Short-circut here.
        short_circuits = [ "exit_to_user_mode_loop", "ret_from_fork" ]
        skip = False
        for s in short_circuits:
            if s in kstack:
                skip = True

        if skip:
            continue

        func_maps = {
                "io_schedule": "disk io",
                "io_schedule_timeout": "disk io",
                "wait_transaction_locked": "disk io",
                "pipe_read": "disk io",
                "pipe_write": "disk io",
                "jbd2_log_wait_commit": "disk io",
                "do_epoll_wait": "disk io",
                "ep_poll": "disk io",
                "wait_woken": "disk io",

                "do_wait": "voluntary wait",
                "do_nanosleep": "voluntary wait",
                "sigsuspend": "voluntary wait",
                "do_sched_yield": "voluntary wait",

                "_cond_resched": "involuntary preempt",
                "__cond_resched": "involuntary preempt",
                "__cond_resched_lock": "involuntary preempt",
                "d_alloc_parallel": "involuntary preempt",
                # This catches vfork(). Not sure this is in the correct bucket
                "wait_for_completion_killable": "involuntary preempt",
                "msleep": "involuntary preempt",
                "wait_for_completion": "involuntary preempt",
                "expand_files": "involuntary preempt",
                "sysvec_reschedule_ipi": "involuntary preempt",

                "futex_wait_queue_me": "userspace locking",

                "rwsem_down_write_slowpath": "kernel locking",
                "rwsem_down_read_slowpath": "kernel locking",
                "__mutex_lock.isra.9": "kernel locking",

                "__skb_wait_for_more_packets": "network io",
                "unix_stream_read_generic": "network io",
                # This is questionable
                "do_select": "network io",

                "page_fault": "page faults",
                "exc_page_fault": "page faults",

        }

        group = None
        for func in func_maps:
            if func in kstack:
                group = func_maps[func]
                break

        if group:
                tasks[pid].times[group] += v.value
                continue

        print("Unknown schedule reason")
        kernel_stack = stack_traces.walk(k.kernel_stack_id)
        for addr in kernel_stack:
            print("    %s" % b.ksym(addr))


    total = {}
    unit = "ms" if args.milliseconds else "s"
    total["abs"] = sum([tasks[t].lifetime for t in tasks])

    calc_time = lambda l: sum([tasks[t].times[g] for t in tasks for g in l])

    io = {}
    io["abs"] = calc_time(["disk io", "network io"])
    io["pct"] = pct(io["abs"], total["abs"])

    uwait = {}
    uwait["abs"] = calc_time(["voluntary wait", "userspace locking"])
    uwait["pct"] = pct(uwait["abs"], total["abs"])
    
    kwait = {}
    kwait["abs"] = calc_time(["involuntary preempt", "kernel locking", "page faults"])
    kwait["pct"] = pct(kwait["abs"], total["abs"])

    oncpu = {}
    oncpu["abs"] = total["abs"] - io["abs"] - uwait["abs"] - kwait["abs"]
    oncpu["pct"] = pct(oncpu["abs"], total["abs"])

    eprint("total: %.2f%s, wall: %.2f%s, on-cpu: %.2f%s (%.1f%%), user wait: %.2f%s (%.1f%%), kernel wait: %.2f%s (%.1f%%), i/o: %.2f%s (%.1f%%)" %
        (time_unit(total["abs"]), unit, time_unit(wall_time), unit,
            time_unit(oncpu["abs"]), unit, oncpu["pct"],
            time_unit(uwait["abs"]), unit, uwait["pct"],
            time_unit(kwait["abs"]), unit, kwait["pct"],
            time_unit(io["abs"]), unit, io["pct"]))

    if not args.verbose:
        sys.exit(0)

    eprint("\ntasks (%d):" % (len(tasks)))
    for task in sort_task_data(tasks, args):
        eprint(task)

    if fh is not sys.stdout:
        fh.close()