Refactor fps_meter so it can be used both as a binary or a library.

BUG=chromium:821046
TEST=manual

Change-Id: I48ab9ce25d1de1defdbe8beee993fddcbe42c8f2
Reviewed-on: https://chromium-review.googlesource.com/988893
Commit-Ready: Cheng-Yu Lee <cylee@google.com>
Tested-by: Cheng-Yu Lee <cylee@google.com>
Reviewed-by: Chung-yih Wang <cywang@chromium.org>
diff --git a/fps_meter/fps_meter.py b/fps_meter/fps_meter.py
index 1e754f8..3826b66 100755
--- a/fps_meter/fps_meter.py
+++ b/fps_meter/fps_meter.py
@@ -2,6 +2,7 @@
 # Copyright 2018 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
+
 """Checks kernel tracing events to get the FPS of a CrOS device.
 
 This script requires root privilege to work properly. It may interfere
@@ -11,22 +12,37 @@
 It doesn't support multi-screen.
 It assumes 60 HZ screen refresh rate.
 
-Sample output:
-# ./fps_meter.py
-trace method: workq
-[111111111111111111111111111111111111111111111111111111111111] FPS: 60
-[111111111111111111111111111111111111111111111111111111111111] FPS: 60
-[11111111111111111111111111111111111111111111111 111111111111] FPS: 59
-[111111111111111111111111111111111111111111111111111111111111] FPS: 60
-[11111111111111 11111111111111111111 11111111 111111111111111] FPS: 57
-[111111111111111111111111111111111    11111111111111111111111] FPS: 56
-[111   111111111111111111111111111111111111111111111111111111] FPS: 57
- ^
- 1 : Frame update count detected in this 1/60 sec interval.
+This script can be used both as a stand alone script or a library.
+
+Sample output (when used as a stand alone script):
+  # ./fps_meter.py
+  trace method: workq
+  [111111111111111111111111111111111111111111111111111111111111] FPS: 60
+  [111111111111111111111111111111111111111111111111111111111111] FPS: 60
+  [11111111111111111111111111111111111111111111111 111111111111] FPS: 59
+  [111111111111111111111111111111111111111111111111111111111111] FPS: 60
+  [11111111111111 11111111111111111111 11111111 111111111111111] FPS: 57
+  [111111111111111111111111111111111    11111111111111111111111] FPS: 56
+  [111   111111111111111111111111111111111111111111111111111111] FPS: 57
+   ^
+     1 : Frame update count detected in this 1/60 sec interval.
+
+Sample Usage (when used as a library).
+  def callback(fps):
+    ...
+
+  with FPSMeter(callback) as fps_meter:
+    ...
+
+When used as a library, it launches two threads to monitor system FPS rate
+periodically. Each time when a FPS rate is sampled, callback() is called with
+the FPS number as its parameter.
 """
 import argparse
 import atexit
 import collections
+import functools
+import logging
 import os
 import re
 import threading
@@ -46,9 +62,11 @@
 TRACE_MARKER = os.path.join(TRACE_ROOT, 'trace_marker')
 REFRESH_RATE = 60
 NOTIFY_STRING = 'notify_collection'
+STOP_STRING = 'stop_tracing'
 
 
 def is_intel_gpu():
+  """Whether is intel GPU."""
   return os.path.isdir('/sys/bus/pci/drivers/i915')
 
 
@@ -78,6 +96,7 @@
 
 
 def write_to_file(content, filename):
+  """Writes content to a file."""
   with open(filename, 'w') as fout:
     fout.write(content)
 
@@ -122,6 +141,7 @@
 
 
 def enable_tracing(trace_method):
+  """Enables tracing."""
   if trace_method == 'workq':
     set_simple_switch('1', WORKQ_SWITCH)
     # There are many workqueue_execute_start events,
@@ -136,31 +156,11 @@
   set_trace_clock()
 
 
-def notify_collection(period_sec=1.0):
-  next_notify_time = time.time() + period_sec
-  while True:
-    # Calling time.sleep(1) may suspend for more than 1 second.
-    # Sleeps until a specific time to avoid error accumulation.
-    sleep_time = next_notify_time - time.time()
-    next_notify_time += period_sec
-    # Skips if current time is larger than next_notify_time.
-    if sleep_time > 0:
-      time.sleep(sleep_time)
-      write_to_file(NOTIFY_STRING, TRACE_MARKER)
-
-
-def start_notify_thread():
-  notify_thread = threading.Thread(target=notify_collection)
-  notify_thread.daemon = True
-  notify_thread.start()
-
-
-def process_trace_buffer(trace_buffer, end_time, verbose=False):
+def get_fps_info(trace_buffer, end_time):
   # Checks all vblanks in the range [end_time - 1.0, end_time].
   frame_info = []
   step = 1.0 / REFRESH_RATE
   step_time = end_time - 1.0 + step
-  fps_count = 0
   frame_times = []
   for _ in range(REFRESH_RATE):
     # Checks if there are vblanks in a refresh interval.
@@ -175,73 +175,142 @@
         frame_info.append('*')
       else:
         frame_info.append(str(step_count))
-      fps_count += 1
     else:
       frame_info.append(' ')
     step_time += step
 
+  return frame_info, frame_times
+
+
+def start_thread(function, args=()):
+  """Starts a thread with given argument."""
+  new_thread = threading.Thread(target=function, args=args)
+  new_thread.start()
+
+
+class FPSMeter():
+  """A FPSMeter to measure system FPS periodically."""
+  def __init__(self, callback):
+    self.is_stopping = threading.Event()
+    self.callback = callback
+
+  def __enter__(self):
+    self.start()
+    return self
+
+  def __exit__(self, type, value, traceback):
+    self.stop()
+
+  def notify_collection(self, period_sec=1.0):
+    """Writes a notification mark periodically."""
+    logging.info('Notification thread is started')
+    next_notify_time = time.time() + period_sec
+    while True:
+      # Calling time.sleep(1) may suspend for more than 1 second.
+      # Sleeps until a specific time to avoid error accumulation.
+      sleep_time = next_notify_time - time.time()
+      next_notify_time += period_sec
+      # Skips if current time is larger than next_notify_time.
+      if sleep_time > 0:
+        if self.is_stopping.wait(sleep_time):
+          logging.info('Leaving notification thread')
+          # So the main loop doesn't stuck in the readline().
+          write_to_file(STOP_STRING, TRACE_MARKER)
+          break
+        write_to_file(NOTIFY_STRING, TRACE_MARKER)
+
+  def main_loop(self, trace_method):
+    """Main loop to parse the trace.
+
+    There are 2 threads involved:
+    Main thread:
+      Using blocking read to get data from trace_pipe.
+    Notify thread: The main thread may wait indifinitely if there
+      is no new trace. Writes to the pipe once per second to avoid
+      the indefinite waiting.
+    """
+    logging.info('Fps meter main thread is started.')
+
+    # Sample trace:
+    # <idle>-0  [000] dNh3 631.905961: drm_vblank_event: crtc=0, seq=65496
+    # <idle>-0  [000] d.h3 631.922681: drm_vblank_event: crtc=0, seq=65497
+    # fps_meter [003] ...1 632.393953: tracing_mark_write: notify_collection
+    # ...
+    re_notify = re.compile(
+        r'.* (\d+.\d+): tracing_mark_write: ' + NOTIFY_STRING)
+    if trace_method == 'workq':
+      re_trace = re.compile(
+          r'.* (\d+.\d+): workqueue_execute_start: work struct ([\da-f]+): '
+          'function intel_atomic_commit_work')
+    elif trace_method == 'flip':
+      re_trace = re.compile(
+          r'.* (\d+.\d+): i915_flip_complete: plane=(\d+), obj=([\da-f]+)')
+    else:
+      re_trace = re.compile(
+          r'.* (\d+.\d+): drm_vblank_event: crtc=(\d+), seq=(\d+)')
+
+    trace_buffer = collections.deque()
+    with open(TRACE_PIPE) as trace_pipe:
+      # The pipe may block a few seconds if using:
+      # for line in trace_pipe
+      while not self.is_stopping.is_set():
+        line = trace_pipe.readline()
+        m_trace = re_trace.match(line)
+        if m_trace:
+          timestamp = float(m_trace.group(1))
+          trace_buffer.append(timestamp)
+        else:
+          m_notify = re_notify.match(line)
+          if m_notify:
+            timestamp = float(m_notify.group(1))
+            self.callback(get_fps_info(trace_buffer, timestamp))
+      logging.info('Leaving fps meter main thread')
+
+  def start(self):
+    """Starts the FPS meter by launching needed threads."""
+    # The notificaiton thread.
+    start_thread(self.notify_collection)
+
+    # The main thread.
+    trace_method = get_trace_method()
+    enable_tracing(trace_method)
+    start_thread(self.main_loop, [trace_method])
+
+  def stop(self):
+    """Stops the FPS meter. Shut down threads."""
+    logging.info('Shutting down FPS meter')
+    self.is_stopping.set()
+
+
+def output_fps_info(verbose, fps_info):
+  frame_info, frame_times = fps_info
+  fps_count = len([f for f in frame_info if f != ' '])
   frame_info_str = ''.join(frame_info)
   print '[%s] FPS: %2d' % (frame_info_str, fps_count)
   if frame_times and verbose:
     print ', '.join(['%.3f' % t for t in frame_times])
 
 
-def main_loop(trace_method, verbose=False):
-  """Main loop to parse the trace.
-
-  There are 2 threads involved:
-  Main thread:
-    Using blocking read to get data from trace_pipe.
-  Notify thread: The main thread may wait indifinitely if there
-    is no new trace. Writes to the pipe once per second to avoid
-    the indefinite waiting.
-  """
-  start_notify_thread()
-
-  # Sample trace:
-  # <idle>-0  [000] dNh3 631.905961: drm_vblank_event: crtc=0, seq=65496
-  # <idle>-0  [000] d.h3 631.922681: drm_vblank_event: crtc=0, seq=65497
-  # fps_meter [003] ...1 632.393953: tracing_mark_write: notify_collection
-  # ...
-  re_notify = re.compile(
-      r'.* (\d+.\d+): tracing_mark_write: ' + NOTIFY_STRING)
-  if trace_method == 'workq':
-    re_trace = re.compile(
-        r'.* (\d+.\d+): workqueue_execute_start: work struct ([\da-f]+): '
-        'function intel_atomic_commit_work')
-  elif trace_method == 'flip':
-    re_trace = re.compile(
-        r'.* (\d+.\d+): i915_flip_complete: plane=(\d+), obj=([\da-f]+)')
-  else:
-    re_trace = re.compile(
-        r'.* (\d+.\d+): drm_vblank_event: crtc=(\d+), seq=(\d+)')
-
-  trace_buffer = collections.deque()
-  with open(TRACE_PIPE) as trace_pipe:
-    # The pipe may block a few seconds if using:
-    # for line in trace_pipe
-    while True:
-      line = trace_pipe.readline()
-      m_trace = re_trace.match(line)
-      if m_trace:
-        timestamp = float(m_trace.group(1))
-        trace_buffer.append(timestamp)
-      else:
-        m_notify = re_notify.match(line)
-        if m_notify:
-          timestamp = float(m_notify.group(1))
-          process_trace_buffer(trace_buffer, timestamp, verbose)
-
-
 if __name__ == '__main__':
   # parsing arguments
   parser = argparse.ArgumentParser(description='Print fps infomation.')
   parser.add_argument('-v', dest='verbose', action='store_true',
                       help='print verbose frame time info')
+  parser.add_argument('-d', dest='debug', action='store_true',
+                      help='print debug message')
   args = parser.parse_args()
 
-  trace_method = get_trace_method()
-  print 'trace method:', trace_method
-  enable_tracing(trace_method)
-  main_loop(trace_method, args.verbose)
+  if args.debug:
+    rootLogger = logging.getLogger()
+    rootLogger.setLevel(logging.DEBUG)
+    # StreamHandler() defaults to stderr.
+    rootLogger.addHandler(logging.StreamHandler())
 
+  print 'FPS meter trace method %s' % get_trace_method()
+  with FPSMeter(functools.partial(output_fps_info, args.verbose)) as fps_meter:
+    while True:
+      try:
+        time.sleep(86400)
+      except KeyboardInterrupt:
+        print 'Exiting...'
+        break