#!/usr/bin/env bpftrace
#include <linux/sched.h>
#include <signal.h>
BEGIN
{
  if ( $# != 1 )
  {
    printf("First argument must be the PID of the process to trace!\n");
    exit();
  }
  @task_state[0]="TASK_RUNNING";
  @task_state[1]="TASK_INTERRUPTIBLE";
  @task_state[2]="TASK_UNINTERRUPTIBLE";
  @task_state[4]="TASK_STOPPED";
  @task_state[8]="TASK_TRACED";
  @task_state[16]="EXIT_ZOMBIE";
  @task_state[32]="EXIT_DEAD";
  @task_state[64]="TASK_DEAD";
  @task_state[128]="TASK_WAKEKILL";
  @task_state[256]="TASK_WAKING";
}
tracepoint:sched:sched_waking
/args->pid == $1/
{
  $time = nsecs;
  printf("[%5u]S:%3u:%10u:%10u:waking\n", args->pid, args->target_cpu, ($time - @leave[args->pid])/1000, 0);
  @waking[args->pid] = $time;
  @current_cpu[args->pid] = args->target_cpu;
}
tracepoint:sched:sched_migrate_task
/args->pid == $1/
{
  printf("[%5u]S:%3u:%10u:%10u:migrate ->%u)\n", args->pid, args->orig_cpu, 0, 0, args->dest_cpu);
  @current_cpu[args->pid] = args->dest_cpu;
}
tracepoint:sched:sched_switch
{
  $pid = $1;
  $time = nsecs;
  // pid is getting scheduled
  if ( args->next_pid == $pid )
  {
    // was this pid woken previously?
    if ( @waking[$pid] != 0 )
    {
      @wake_latency[$pid] = $time - @waking[$pid];
      @waking[$pid] = 0;
      printf("[%5u]S:%3u:%10u:%10u:on cpu\n", $pid, @current_cpu[$pid], @wake_latency[$pid]/1000, 0);
    }
    // otherwise the pid is pushed off cpu involuntarily, and is getting back on cpu:
    else
    {
      @reschedule[$pid] = $time - @leave[$pid];
      printf("[%5u]S:%3u:%10u:%10u:on cpu\n", $pid, @current_cpu[$pid], @reschedule[$pid]/1000, 0);
    }
    @schedule[$pid] = $time;
  }
  // pid is getting scheduled off cpu
  if ( args->prev_pid == $pid )
  {
    @leave[$pid] = $time;
    // did we get on cpu after getting woken?
    if ( @wake_latency[$pid] != 0 )
    {
      printf("[%5u]S:%3u:%10u:%10u:off cpu:%s\n", $pid, @current_cpu[$pid], @wake_latency[$pid]/1000, ($time - @schedule[$pid])/1000, @task_state[args->prev_state]);
    }
    // otherwise the pid was still running and pushed off cpu, and now gotten back onto the cpu
    else
    {
      printf("[%5u]S:%3u:%10u:%10u:off cpu:%s\n", $pid, @current_cpu[$pid], @reschedule[$pid]/1000, ($time - @schedule[$pid])/1000, @task_state[args->prev_state]);
    }
  }
}
END
{
  clear(@task_state);
  clear(@waking);
  clear(@wake_latency);
  clear(@reschedule);
  clear(@leave);
  clear(@schedule);
  clear(@current_cpu);
}