add contextmanager'd exiters and change return type of run()

src/python/pants/bin/daemon_pants_runner.py

@@ -237,6 +237,27 @@ def nailgunned_stdio(cls, sock, env, handle_stdin=True):
       ) as finalizer:
         yield finalizer
 
+  @contextmanager
+  def daemon_exiter(cls, exiter):
+    """A contextmanager which temporarily overrides the exiter."""
+    try:
+      cls._exiter = exiter
+      yield exiter
+    finally:
+      pass
+
+  @contextmanager
+  def set_exiter(cls, exiter):
+    """A contextmanager which temporarily overrides the exiter."""
+    try:
+      previous_exiter = cls._exiter
+      cls._exiter = exiter
+      ExceptionSink.reset_exiter(exiter)
+      yield cls._exiter
+    finally:
+      cls._exiter = previous_exiter
+      ExceptionSink.reset_exiter(previous_exiter)
+
   # TODO: there's no testing for this method, and this caused a user-visible failure -- see #7008!
   def _raise_deferred_exc(self):
     """Raises deferred exceptions from the daemon's synchronous path in the post-fork client."""
@@ -249,10 +270,6 @@ def _raise_deferred_exc(self):
         # destructuring), treat it like a bare exception.
         raise self._deferred_exception
 
-  def _maybe_get_client_start_time_from_env(self, env):
-    client_start_time = env.pop('PANTSD_RUNTRACKER_CLIENT_START_TIME', None)
-    return None if client_start_time is None else float(client_start_time)
-
   def run(self):
     """Fork, daemonize and invoke self.post_fork_child() (via ProcessManager).
 
@@ -307,8 +324,10 @@ def post_fork_child(self):
       service.terminate()
 
     # Invoke a Pants run with stdio redirected and a proxied environment.
+    daemonExiter = DaemonExiter(self._socket)
     with self.nailgunned_stdio(self._socket, self._env) as finalizer,\
-         hermetic_environment_as(**self._env):
+         hermetic_environment_as(**self._env),\
+         self.set_exiter(daemonExiter) as daemonExiter:
       try:
         # Setup the Exiter's finalizer.
         self._exiter.set_finalizer(finalizer)
@@ -318,7 +337,6 @@ def post_fork_child(self):
 
         # Re-raise any deferred exceptions, if present.
         self._raise_deferred_exc()
-
         # Otherwise, conduct a normal run.
         runner = LocalPantsRunner.create(
           self._exiter,
@@ -328,19 +346,19 @@ def post_fork_child(self):
           self._graph_helper,
           self._options_bootstrapper
         )
-        runner.set_start_time(self._maybe_get_client_start_time_from_env(self._env))
-        runner.run()
+        exit_code = runner.run()
+        daemonExiter.exit(exit_code)
       except KeyboardInterrupt:
-        self._exiter.exit_and_fail('Interrupted by user.\n')
+        daemonExiter.exit_and_fail('Interrupted by user.\n')
       except GracefulTerminationException as e:
         ExceptionSink.log_exception(
           'Encountered graceful termination exception {}; exiting'.format(e))
-        self._exiter.exit(e.exit_code)
+        daemonExiter.exit(e.exit_code)
       except Exception:
         # TODO: We override sys.excepthook above when we call ExceptionSink.set_exiter(). That
         # excepthook catches `SignalHandledNonLocalExit`s from signal handlers, which isn't
         # happening here, so something is probably overriding the excepthook. By catching Exception
         # and calling this method, we emulate the normal, expected sys.excepthook override.
         ExceptionSink._log_unhandled_exception_and_exit()
       else:
-        self._exiter.exit(PANTS_SUCCEEDED_EXIT_CODE)
+        daemonExiter.exit(PANTS_SUCCEEDED_EXIT_CODE)

src/python/pants/bin/local_pants_runner.py

@@ -5,7 +5,10 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import logging
+import time
 import traceback
+from contextlib import contextmanager
+
 from builtins import object, str
 
 from pants.base.build_environment import get_buildroot
@@ -82,6 +85,17 @@ def exit(self, result=PANTS_SUCCEEDED_EXIT_CODE, msg=None, *args, **kwargs):
 class LocalPantsRunner(object):
   """Handles a single pants invocation running in the process-local context."""
 
+  @contextmanager
+  def set_exiter(self, exiter):
+    """A contextmanager which temporarily overrides the exiter."""
+    previous_exiter = self._exiter if self._exiter else exiter
+    try:
+      ExceptionSink.reset_exiter(exiter)
+      yield
+    finally:
+      self._exiter = previous_exiter
+      ExceptionSink.reset_exiter(self._exiter)
+
   @staticmethod
   def parse_options(args, env, setup_logging=False, options_bootstrapper=None):
     options_bootstrapper = options_bootstrapper or OptionsBootstrapper.create(args=args, env=env)
@@ -214,7 +228,7 @@ def __init__(self, build_root, exiter, options, options_bootstrapper, build_conf
     self._repro = None
     self._global_options = options.for_global_scope()
 
-  def set_start_time(self, start_time):
+  def _set_start_time(self, start_time):
     # Launch RunTracker as early as possible (before .run() is called).
     self._run_tracker = RunTracker.global_instance()
     self._reporting = Reporting.global_instance()
@@ -227,13 +241,9 @@ def set_start_time(self, start_time):
     if self._repro:
       self._repro.capture(self._run_tracker.run_info.get_as_dict())
 
-    # The __call__ method of the Exiter allows for the prototype pattern.
-    self._exiter = LocalExiter(self._run_tracker, self._repro, exiter=self._exiter)
-    ExceptionSink.reset_exiter(self._exiter)
-
   def run(self):
     with maybe_profiled(self._profile_path):
-      self._run()
+      return self._run()
 
   def _maybe_run_v1(self):
     if not self._global_options.v1:
@@ -289,20 +299,25 @@ def _compute_final_exit_code(*codes):
     return max_code
 
   def _run(self):
-    try:
-      engine_result = self._maybe_run_v2()
-      goal_runner_result = self._maybe_run_v1()
-    finally:
+    self._set_start_time(time.time())
+
+    # wrap the outer exiter
+    local_exiter = LocalExiter(self._run_tracker, self._repro, exiter=self._exiter)
+    with self.set_exiter(local_exiter):
       try:
-        run_tracker_result = self._run_tracker.end()
-      except ValueError as e:
-        # Calling .end() sometimes writes to a closed file, so we return a dummy result here.
-        logger.exception(e)
-        run_tracker_result = PANTS_SUCCEEDED_EXIT_CODE
-
-    final_exit_code = self._compute_final_exit_code(
-      engine_result,
-      goal_runner_result,
-      run_tracker_result
-    )
-    self._exiter.exit(final_exit_code)
+        engine_result = self._maybe_run_v2()
+        goal_runner_result = self._maybe_run_v1()
+      finally:
+        try:
+          run_tracker_result = self._run_tracker.end()
+        except ValueError as e:
+          # Calling .end() sometimes writes to a closed file, so we return a dummy result here.
+          logger.exception(e)
+          run_tracker_result = PANTS_SUCCEEDED_EXIT_CODE
+
+      final_exit_code = self._compute_final_exit_code(
+        engine_result,
+        goal_runner_result,
+        run_tracker_result
+      )
+      return final_exit_code

src/python/pants/bin/pants_runner.py

@@ -72,5 +72,4 @@ def run(self):
         self._env,
         options_bootstrapper=options_bootstrapper
     )
-    runner.set_start_time(self._start_time)
     return runner.run()