text: #!/usr/bin/env python
  #
  # sema    Time semaphore actions in PostgreSQL and print
  #         wait/hold time as a histogram. For Linux, uses BCC, eBPF.
  #
  # usage: sema $PG_BIN/postgres [-p PID] [-d]
  
  
  from __future__ import print_function
  from time import sleep
  from bcc import BPF
  
  import argparse
  import ctypes as ct
  import signal
  
  
  text = """
  #include <linux/ptrace.h>
  
  struct pg_semaphore {
      u32 pid;
      u64 acquired;
      u64 released;
  };
  
  struct message {
      char text[100];
  };
  
  BPF_PERF_OUTPUT(events);
  BPF_PERF_OUTPUT(messages);
  
  BPF_HASH(lock_hold, u32, struct pg_semaphore);
  BPF_HASH(lock_wait, u32, struct pg_semaphore);
  
  // Histogram of semaphore hold times
  BPF_HISTOGRAM(lock_hold_semaphore_hist, u64);
  
  // Histogram of semaphore wait times
  BPF_HISTOGRAM(lock_wait_semaphore_hist, u64);
  
  void probe_semaphore_lock_start(struct pt_regs *ctx)
  {
      u64 now = bpf_ktime_get_ns();
      u32 pid = bpf_get_current_pid_tgid();
      struct pg_semaphore data = {};
      data.acquired = now;
      data.pid = pid;
      data.released = 0;
  
      events.perf_submit(ctx, &data, sizeof(data));
      struct pg_semaphore *test = lock_wait.lookup(&pid);
      if (test != 0)
      {
          struct message msg = {};
          strcpy(msg.text, "Semaphore is overwritten");
          messages.perf_submit(ctx, &msg, sizeof(msg));
      }
  
      lock_wait.update(&pid, &data);
  }
  
  void probe_semaphore_lock_finish(struct pt_regs *ctx)
  {
      u64 now = bpf_ktime_get_ns();
      u32 pid = bpf_get_current_pid_tgid();
      struct pg_semaphore data = {};
      struct pg_semaphore *wait_data = lock_wait.lookup(&pid);
  
      if (wait_data != 0)
      {
          u64 wait_time = now - wait_data->acquired;
          u64 semaphore_slot = bpf_log2l(wait_time / 1000);
          lock_wait_semaphore_hist.increment(semaphore_slot);
          lock_wait.delete(&pid);
      }
  
      data.acquired = now;
      data.pid = pid;
      data.released = 0;
  
      events.perf_submit(ctx, &data, sizeof(data));
      struct pg_semaphore *test = lock_hold.lookup(&pid);
      if (test != 0)
      {
          struct message msg = {};
          strcpy(msg.text, "Semaphore is overwritten");
          messages.perf_submit(ctx, &msg, sizeof(msg));
      }
  
      lock_hold.update(&pid, &data);
  }
  
  void probe_semaphore_unlock(struct pt_regs *ctx)
  {
      u64 now = bpf_ktime_get_ns();
      u32 pid = bpf_get_current_pid_tgid();
      struct pg_semaphore *data = lock_hold.lookup(&pid);
      if (data == 0)
      {
          struct message msg = {};
          strcpy(msg.text, "Semaphore is missing");
          messages.perf_submit(ctx, &msg, sizeof(msg));
          return;
      }
  
      u64 hold_time = now - data->acquired;
      data->released = now;
  
      events.perf_submit(ctx, data, sizeof(*data));
      u64 semaphore_slot = bpf_log2l(hold_time / 1000);
      lock_hold_semaphore_hist.increment(semaphore_slot);
      lock_hold.delete(&pid);
  }
  """
  
  
  def attach(bpf, binary_path, pid=-1):
      bpf.attach_uprobe(
          name=binary_path,
          sym="PGSemaphoreLock",
          fn_name="probe_semaphore_lock_start",
          pid=pid)
      bpf.attach_uretprobe(
          name=binary_path,
          sym="PGSemaphoreLock",
          fn_name="probe_semaphore_lock_finish",
          pid=pid)
  
      bpf.attach_uprobe(
          name=binary_path,
          sym="PGSemaphoreUnlock",
          fn_name="probe_semaphore_unlock",
          pid=pid)
  
  
  # signal handler
  def signal_ignore(signal, frame):
      print()
  
  
  class Data(ct.Structure):
      _fields_ = [("pid", ct.c_ulong),
                  ("mode", ct.c_int),
                  ("acquired", ct.c_ulonglong),
                  ("released", ct.c_ulonglong)]
  
  
  class Message(ct.Structure):
      _fields_ = [("text", ct.c_char * 100)]
  
  
  def print_messages(cpu, data, size):
      msg = ct.cast(data, ct.POINTER(Message)).contents
      print(msg.text)
  
  
  def print_event(cpu, data, size):
      event = ct.cast(data, ct.POINTER(Data)).contents
      print("Event: acquired {} released {} pid {}".format(
          event.acquired, event.released, event.pid))
  
  
  def run(args):
      print("Attaching...")
      bpf = BPF(text=text)
      attach(bpf, args.path, args.pid)
      lock_hold_semaphore_hist = bpf["lock_hold_semaphore_hist"]
      lock_wait_semaphore_hist = bpf["lock_wait_semaphore_hist"]
      exiting = False
  
      if args.debug:
          bpf["events"].open_perf_buffer(print_event)
          bpf["messages"].open_perf_buffer(print_messages)
  
      print("Listening...")
      while True:
          try:
              sleep(1)
              if args.debug:
                  bpf.perf_buffer_poll()
          except KeyboardInterrupt:
              exiting = True
              # as cleanup can take many seconds, trap Ctrl-C:
              signal.signal(signal.SIGINT, signal_ignore)
  
          if exiting:
              print("Detaching...")
              break
  
      print("Semaphore holding time")
      lock_hold_semaphore_hist.print_log2_hist("hold time (us)")
      print("")
  
      print("Semaphore waiting time")
      lock_wait_semaphore_hist.print_log2_hist("wait time (us)")
      print("")
  
  
  def parse_args():
      parser = argparse.ArgumentParser(
          description="Time LWLocks in PostgreSQL",
          formatter_class=argparse.RawDescriptionHelpFormatter)
      parser.add_argument("path", type=str, help="path to target binary")
      parser.add_argument(
          "-p", "--pid", type=int, default=-1,
          help="trace this PID only")
      parser.add_argument(
          "-d", "--debug", action='store_true', default=False,
          help="debug mode")
  
      return parser.parse_args()
  
  
  if __name__ == "__main__":
      run(parse_args())
  