metrics: Separate metrics collection and print notice logic.

As it is right now, when a function is collecting metrics it prints
the notice and calls sys.exit() to ensure the notice is the last thing
that is printed.

git-cl split has to call git-cl upload multiple times, but once it has
been called once, it exits, so only the first cl is uploaded.

This separates metrics collection from notice printing, so that the
function that is collecting metrics behaves like a function that isn't.
It also makes sure we don't collect metrics for multiple functions at
the same time.

Bug: 868280
Change-Id: Ic58ebe7d19e09ed85fa8b0af76dcbf608ee4c9bc
Reviewed-on: https://chromium-review.googlesource.com/1153503
Reviewed-by: Aaron Gable <agable@chromium.org>
Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>

gclient.py

@@ -3026,10 +3026,7 @@ def main(argv):
 
 
 if '__main__' == __name__:
-  try:
+  with metrics.collector.print_notice_and_exit():
     sys.exit(main(sys.argv[1:]))
-  except KeyboardInterrupt:
-    sys.stderr.write('interrupted\n')
-    sys.exit(1)
 
 # vim: ts=2:sw=2:tw=80:et:

git_cl.py

@@ -6083,8 +6083,5 @@ if __name__ == '__main__':
   # unit testing.
   fix_encoding.fix_encoding()
   setup_color.init()
-  try:
+  with metrics.collector.print_notice_and_exit():
     sys.exit(main(sys.argv[1:]))
-  except KeyboardInterrupt:
-    sys.stderr.write('interrupted\n')
-    sys.exit(1)

git_nav_downstream.py

@@ -62,8 +62,5 @@ def main(args):
 
 
 if __name__ == '__main__':
-  try:
+  with metrics.collector.print_notice_and_exit():
     sys.exit(main(sys.argv[1:]))
-  except KeyboardInterrupt:
-    sys.stderr.write('interrupted\n')
-    sys.exit(1)

metrics.py

@@ -3,6 +3,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import contextlib
 import functools
 import json
 import os
@@ -107,6 +108,7 @@ class MetricsCollector(object):
     self._reported_metrics = {}
     self._config = _Config()
     self._collecting_metrics = False
+    self._collect_custom_metrics = True
 
   @property
   def config(self):
@@ -117,8 +119,18 @@ class MetricsCollector(object):
     return self._collecting_metrics
 
   def add(self, name, value):
-    with self._metrics_lock:
-      self._reported_metrics[name] = value
+    if self._collect_custom_metrics:
+      with self._metrics_lock:
+        self._reported_metrics[name] = value
+
+  @contextlib.contextmanager
+  def pause_metrics_collection(self):
+    collect_custom_metrics = self._collect_custom_metrics
+    self._collect_custom_metrics = False
+    try:
+      yield
+    finally:
+      self._collect_custom_metrics = collect_custom_metrics
 
   def _upload_metrics_data(self):
     """Upload the metrics data to the AppEngine app."""
@@ -129,19 +141,20 @@ class MetricsCollector(object):
     p.stdin.write(json.dumps(self._reported_metrics))
 
   def _collect_metrics(self, func, command_name, *args, **kwargs):
-    # If the user hasn't opted in or out, and the countdown is not yet 0, just
-    # display the notice.
-    if self.config.opted_in == None and self.config.countdown > 0:
-      metrics_utils.print_notice(self.config.countdown)
-      self.config.decrease_countdown()
-      func(*args, **kwargs)
-      return
+    # If we're already collecting metrics, just execute the function.
+    # e.g. git-cl split invokes git-cl upload several times to upload each
+    # splitted CL.
+    if self.collecting_metrics:
+      # Don't collect metrics for this function.
+      # e.g. Don't record the arguments git-cl split passes to git-cl upload.
+      with self.pause_metrics_collection():
+        return func(*args, **kwargs)
 
     self._collecting_metrics = True
     self.add('command', command_name)
     try:
       start = time.time()
-      func(*args, **kwargs)
+      result = func(*args, **kwargs)
       exception = None
     # pylint: disable=bare-except
     except:
@@ -149,19 +162,9 @@ class MetricsCollector(object):
     finally:
       self.add('execution_time', time.time() - start)
 
-    # Print the exception before the metrics notice, so that the notice is
-    # clearly visible even if gclient fails.
-    if exception and not isinstance(exception[1], SystemExit):
-      traceback.print_exception(*exception)
-
     exit_code = metrics_utils.return_code_from_exception(exception)
     self.add('exit_code', exit_code)
 
-    # Print the metrics notice only if the user has not explicitly opted in
-    # or out.
-    if self.config.opted_in is None:
-      metrics_utils.print_notice(self.config.countdown)
-
     # Add metrics regarding environment information.
     self.add('timestamp', metrics_utils.seconds_to_weeks(time.time()))
     self.add('python_version', metrics_utils.get_python_version())
@@ -170,31 +173,66 @@ class MetricsCollector(object):
     self.add('depot_tools_age', metrics_utils.get_repo_timestamp(DEPOT_TOOLS))
 
     self._upload_metrics_data()
-    sys.exit(exit_code)
+    if exception:
+      raise exception[0], exception[1], exception[2]
+    return result
 
   def collect_metrics(self, command_name):
     """A decorator used to collect metrics over the life of a function.
 
     This decorator executes the function and collects metrics about the system
-    environment and the function performance. It also catches all the Exceptions
-    and invokes sys.exit once the function is done executing.
+    environment and the function performance.
     """
     def _decorator(func):
       # Do this first so we don't have to read, and possibly create a config
       # file.
       if DISABLE_METRICS_COLLECTION:
         return func
-      # If the user has opted out or the user is not a googler, then there is no
-      # need to do anything.
-      if self.config.opted_in == False or not self.config.is_googler:
+      # Don't collect the metrics unless the user is a googler, the user has
+      # opted in, or the countdown has expired.
+      if (not self.config.is_googler or self.config.opted_in == False
+          or (self.config.opted_in is None and self.config.countdown > 0)):
         return func
       # Otherwise, collect the metrics.
       # Needed to preserve the __name__ and __doc__ attributes of func.
       @functools.wraps(func)
       def _inner(*args, **kwargs):
-        self._collect_metrics(func, command_name, *args, **kwargs)
+        return self._collect_metrics(func, command_name, *args, **kwargs)
       return _inner
     return _decorator
 
+  @contextlib.contextmanager
+  def print_notice_and_exit(self):
+    """A context manager used to print the notice and terminate execution.
+
+    This decorator executes the function and prints the monitoring notice if
+    necessary. If an exception is raised, we will catch it, and print it before
+    printing the metrics collection notice.
+    This will call sys.exit() with an appropriate exit code to ensure the notice
+    is the last thing printed."""
+    # Needed to preserve the __name__ and __doc__ attributes of func.
+    try:
+      yield
+      exception = None
+    # pylint: disable=bare-except
+    except:
+      exception = sys.exc_info()
+
+    # Print the exception before the metrics notice, so that the notice is
+    # clearly visible even if gclient fails.
+    if exception:
+      if isinstance(exception[1], KeyboardInterrupt):
+        sys.stderr.write('Interrupted\n')
+      elif not isinstance(exception[1], SystemExit):
+        traceback.print_exception(*exception)
+
+    # Print the notice
+    if (not DISABLE_METRICS_COLLECTION and self.config.is_googler
+        and self.config.opted_in is None):
+      metrics_utils.print_notice(self.config.countdown)
+      self.config.decrease_countdown()
+
+    sys.exit(metrics_utils.return_code_from_exception(exception))
+
 
 collector = MetricsCollector()

tests/metrics_test.py

@@ -84,7 +84,6 @@ class MetricsCollectorTest(unittest.TestCase):
     self.assertTrue(self.collector.collecting_metrics)
     self.assertEqual(collected_metrics, expected_metrics)
 
-
   def test_collects_system_information(self):
     """Tests that we collect information about the runtime environment."""
     self.FileRead.side_effect = [
@@ -94,9 +93,7 @@ class MetricsCollectorTest(unittest.TestCase):
     def fun():
       pass
 
-    with self.assertRaises(SystemExit) as cm:
-      fun()
-    self.assertEqual(cm.exception.code, 0)
+    fun()
     self.assert_collects_metrics()
 
   def test_collects_added_metrics(self):
@@ -108,9 +105,7 @@ class MetricsCollectorTest(unittest.TestCase):
     def fun():
       self.collector.add('foo', 'bar')
 
-    with self.assertRaises(SystemExit) as cm:
-      fun()
-    self.assertEqual(cm.exception.code, 0)
+    fun()
     self.assert_collects_metrics({'foo': 'bar'})
 
   def test_collects_metrics_when_opted_in(self):
@@ -122,9 +117,7 @@ class MetricsCollectorTest(unittest.TestCase):
     def fun():
       pass
 
-    with self.assertRaises(SystemExit) as cm:
-      fun()
-    self.assertEqual(cm.exception.code, 0)
+    fun()
     self.assert_collects_metrics()
 
   @mock.patch('metrics.DISABLE_METRICS_COLLECTION', True)
@@ -192,47 +185,87 @@ class MetricsCollectorTest(unittest.TestCase):
     self.assertFalse(self.collector.collecting_metrics)
     self.assertTrue(self.collector.config.is_googler)
     self.assertFalse(self.collector.config.opted_in)
-    # The countdown should've decreased after the invocation.
-    self.assertEqual(self.collector.config.countdown, 0)
+    self.assertEqual(self.collector.config.countdown, 1)
     # Assert that we did not try to upload any metrics.
     self.assertFalse(self.Popen.called)
 
-  def test_prints_notice_non_zero_countdown(self):
-    """Tests that a notice is printed while the countdown is non-zero."""
+  def test_handles_exceptions(self):
+    """Tests that exception are caught and we exit with an appropriate code."""
     self.FileRead.side_effect = [
-        '{"is-googler": true, "countdown": 1234, "opt-in": null}'
+        '{"is-googler": true, "countdown": 0, "opt-in": true}'
     ]
     @self.collector.collect_metrics('fun')
     def fun():
-      pass
-    fun()
-    self.print_notice.assert_called_once_with(1234)
+      raise ValueError
 
-  def test_prints_notice_zero_countdown(self):
-    """Tests that a notice is printed when the countdown reaches 0."""
+    # When an exception is raised, we should catch it, update exit-code,
+    # collect metrics, and re-raise it.
+    with self.assertRaises(ValueError):
+      fun()
+    self.assert_collects_metrics({'exit_code': 1})
+
+  def test_handles_system_exit(self):
+    """Tests that the sys.exit code is respected and metrics are collected."""
     self.FileRead.side_effect = [
-        '{"is-googler": true, "countdown": 0, "opt-in": null}'
+        '{"is-googler": true, "countdown": 0, "opt-in": true}'
     ]
     @self.collector.collect_metrics('fun')
     def fun():
-      pass
+      sys.exit(0)
 
+    # When an exception is raised, we should catch it, update exit-code,
+    # collect metrics, and re-raise it.
     with self.assertRaises(SystemExit) as cm:
       fun()
     self.assertEqual(cm.exception.code, 0)
-    self.print_notice.assert_called_once_with(0)
+    self.assert_collects_metrics({'exit_code': 0})
 
-  def test_doesnt_print_notice_opted_in(self):
-    """Tests that a notice is not printed when the user opts-in."""
+  def test_handles_system_exit_non_zero(self):
+    """Tests that the sys.exit code is respected and metrics are collected."""
     self.FileRead.side_effect = [
         '{"is-googler": true, "countdown": 0, "opt-in": true}'
     ]
     @self.collector.collect_metrics('fun')
     def fun():
-      pass
+      sys.exit(123)
 
+    # When an exception is raised, we should catch it, update exit-code,
+    # collect metrics, and re-raise it.
     with self.assertRaises(SystemExit) as cm:
       fun()
+    self.assertEqual(cm.exception.code, 123)
+    self.assert_collects_metrics({'exit_code': 123})
+
+  def test_prints_notice_non_zero_countdown(self):
+    """Tests that a notice is printed while the countdown is non-zero."""
+    self.FileRead.side_effect = [
+        '{"is-googler": true, "countdown": 1234, "opt-in": null}'
+    ]
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        pass
+    self.assertEqual(cm.exception.code, 0)
+    self.print_notice.assert_called_once_with(1234)
+
+  def test_prints_notice_zero_countdown(self):
+    """Tests that a notice is printed when the countdown reaches 0."""
+    self.FileRead.side_effect = [
+        '{"is-googler": true, "countdown": 0, "opt-in": null}'
+    ]
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        pass
+    self.assertEqual(cm.exception.code, 0)
+    self.print_notice.assert_called_once_with(0)
+
+  def test_doesnt_print_notice_opted_in(self):
+    """Tests that a notice is not printed when the user opts-in."""
+    self.FileRead.side_effect = [
+        '{"is-googler": true, "countdown": 0, "opt-in": true}'
+    ]
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        pass
     self.assertEqual(cm.exception.code, 0)
     self.assertFalse(self.print_notice.called)
 
@@ -241,60 +274,58 @@ class MetricsCollectorTest(unittest.TestCase):
     self.FileRead.side_effect = [
         '{"is-googler": true, "countdown": 0, "opt-in": false}'
     ]
-    @self.collector.collect_metrics('fun')
-    def fun():
-      pass
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        pass
+    self.assertEqual(cm.exception.code, 0)
+    self.assertFalse(self.print_notice.called)
 
-    fun()
+  @mock.patch('metrics.DISABLE_METRICS_COLLECTION', True)
+  def test_doesnt_print_notice_disable_metrics_collection(self):
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        pass
+    self.assertEqual(cm.exception.code, 0)
     self.assertFalse(self.print_notice.called)
+    # We shouldn't have tried to read the config file.
+    self.assertFalse(self.FileRead.called)
 
-  def test_handles_exceptions(self):
+  def test_print_notice_handles_exceptions(self):
     """Tests that exception are caught and we exit with an appropriate code."""
     self.FileRead.side_effect = [
-        '{"is-googler": true, "countdown": 0, "opt-in": true}'
+        '{"is-googler": true, "countdown": 0, "opt-in": null}'
     ]
-    @self.collector.collect_metrics('fun')
-    def fun():
-      raise ValueError
-
-    # When an exception is raised, we should catch it, print the traceback and
-    # invoke sys.exit with a non-zero exit code.
+    # print_notice should catch the exception, print it and invoke sys.exit()
     with self.assertRaises(SystemExit) as cm:
-      fun()
+      with self.collector.print_notice_and_exit():
+        raise ValueError
     self.assertEqual(cm.exception.code, 1)
-    self.assert_collects_metrics({'exit_code': 1})
+    self.assertTrue(self.print_notice.called)
 
-  def test_handles_system_exit(self):
-    """Tests that the sys.exit code is respected and metrics are collected."""
+  def test_print_notice_handles_system_exit(self):
+    """Tests that the sys.exit code is respected and a notice is displayed."""
     self.FileRead.side_effect = [
-        '{"is-googler": true, "countdown": 0, "opt-in": true}'
+        '{"is-googler": true, "countdown": 0, "opt-in": null}'
     ]
-    @self.collector.collect_metrics('fun')
-    def fun():
-      sys.exit(0)
-
-    # When an exception is raised, we should catch it, print the traceback and
-    # invoke sys.exit with a non-zero exit code.
+    # print_notice should catch the exception, print it and invoke sys.exit()
     with self.assertRaises(SystemExit) as cm:
-      fun()
+      with self.collector.print_notice_and_exit():
+        sys.exit(0)
     self.assertEqual(cm.exception.code, 0)
-    self.assert_collects_metrics({'exit_code': 0})
+    self.assertTrue(self.print_notice.called)
 
-  def test_handles_system_exit_non_zero(self):
-    """Tests that the sys.exit code is respected and metrics are collected."""
+  def test_print_notice_handles_system_exit_non_zero(self):
+    """Tests that the sys.exit code is respected and a notice is displayed."""
     self.FileRead.side_effect = [
-        '{"is-googler": true, "countdown": 0, "opt-in": true}'
+        '{"is-googler": true, "countdown": 0, "opt-in": null}'
     ]
-    @self.collector.collect_metrics('fun')
-    def fun():
-      sys.exit(123)
-
-    # When an exception is raised, we should catch it, print the traceback and
-    # invoke sys.exit with a non-zero exit code.
+    # When an exception is raised, we should catch it, update exit-code,
+    # collect metrics, and re-raise it.
     with self.assertRaises(SystemExit) as cm:
-      fun()
+      with self.collector.print_notice_and_exit():
+        sys.exit(123)
     self.assertEqual(cm.exception.code, 123)
-    self.assert_collects_metrics({'exit_code': 123})
+    self.assertTrue(self.print_notice.called)
 
   def test_counts_down(self):
     """Tests that the countdown works correctly."""
@@ -303,18 +334,23 @@ class MetricsCollectorTest(unittest.TestCase):
     ]
 
     # We define multiple functions to ensure it has no impact on countdown.
+    @self.collector.collect_metrics('barn')
+    def _barn():
+      pass
     @self.collector.collect_metrics('fun')
-    def fun():
+    def _fun():
       pass
-    @self.collector.collect_metrics('foon')
-    def _foon():
+    def foo_main():
       pass
 
     # Assert that the countdown hasn't decrease yet.
     self.assertFalse(self.FileWrite.called)
     self.assertEqual(self.collector.config.countdown, 10)
 
-    fun()
+    with self.assertRaises(SystemExit) as cm:
+      with self.collector.print_notice_and_exit():
+        foo_main()
+    self.assertEqual(cm.exception.code, 0)
 
     # Assert that the countdown decreased by one, and the config file was
     # updated.
@@ -328,6 +364,27 @@ class MetricsCollectorTest(unittest.TestCase):
     self.assertEqual(json.loads(config),
                      {'is-googler': True, 'countdown': 9, 'opt-in': None})
 
+  def test_nested_functions(self):
+    """Tests that a function can call another function for which metrics are
+    collected."""
+    self.FileRead.side_effect = [
+        '{"is-googler": true, "countdown": 0, "opt-in": true}'
+    ]
+
+    @self.collector.collect_metrics('barn')
+    def barn():
+      self.collector.add('barn-metric', 1)
+      return 1000
+    @self.collector.collect_metrics('fun')
+    def fun():
+      result = barn()
+      self.collector.add('fun-metric', result + 1)
+
+    fun()
+
+    # Assert that we collected metrics for fun, but not for barn.
+    self.assert_collects_metrics({'fun-metric': 1001})
+
 
 if __name__ == '__main__':
   unittest.main()