From b2d363271681e3f6bf78eb43636b921d54d737c7 Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Fri, 13 Nov 2020 06:24:31 -0500
Subject: [PATCH] tools/biolatpcts: Support measuring overall latencies between
 two events

* Interpret negative interval as infinite and output once as soon as init is
  complete to indicate readiness.

* Output on SIGUSR1.

* When exit is requested through SIGINT, TERM or HUP, output the latency
  distribution one more time before exiting.

Combined, this allows using biolatpcts to measure the overall latency
distribution between two events.
---
 tools/biolatpcts.py | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/tools/biolatpcts.py b/tools/biolatpcts.py
index 7776a6f5e147..1b963dab1c47 100755
--- a/tools/biolatpcts.py
+++ b/tools/biolatpcts.py
@@ -15,21 +15,31 @@
 from __future__ import print_function
 from bcc import BPF
 from time import sleep
+from threading import Event
 import argparse
 import json
 import sys
 import os
+import signal
 
 description = """
 Monitor IO latency distribution of a block device
 """
 
-parser = argparse.ArgumentParser(description = description,
+epilog = """
+When interval is infinite, biolatpcts will print out result once the
+initialization is complete to indicate readiness. Once initialized,
+biolatpcts will output whenever it receives SIGUSR1 and before exiting on
+SIGINT, SIGTERM or SIGHUP. This can be used to obtain latency distribution
+between two events.
+"""
+
+parser = argparse.ArgumentParser(description = description, epilog = epilog,
                                  formatter_class = argparse.ArgumentDefaultsHelpFormatter)
 parser.add_argument('dev', metavar='DEV', type=str,
                     help='Target block device (/dev/DEVNAME, DEVNAME or MAJ:MIN)')
 parser.add_argument('-i', '--interval', type=int, default=3,
-                    help='Report interval')
+                    help='Report interval (0: exit after startup, -1: infinite)')
 parser.add_argument('-w', '--which', choices=['from-rq-alloc', 'after-rq-alloc', 'on-device'],
                     default='on-device', help='Which latency to measure')
 parser.add_argument('-p', '--pcts', metavar='PCT,...', type=str,
@@ -206,8 +216,28 @@ def format_usec(lat):
 if args.interval == 0:
     sys.exit(0)
 
-while True:
-    sleep(args.interval)
+# Set up signal handling so that we print the result on USR1 and before
+# exiting on a signal. Combined with infinite interval, this can be used to
+# obtain overall latency distribution between two events.
+keep_running = True
+result_req = Event()
+def sig_handler(sig, frame):
+    global keep_running, result_req
+    if sig != signal.SIGUSR1:
+        keep_running = False
+    result_req.set()
+
+for sig in (signal.SIGUSR1, signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
+    signal.signal(sig, sig_handler)
+
+# If infinite interval, always trigger the first output so that the caller
+# can tell when initialization is complete.
+if args.interval < 0:
+    result_req.set();
+
+while keep_running:
+    result_req.wait(args.interval if args.interval > 0 else None)
+    result_req.clear()
 
     rwdf_total = [0] * 4;