add lifecycle handlers to allow extensions to hook session lifecycle

src/python/pants/bin/local_pants_runner.py

@@ -4,7 +4,7 @@
 import logging
 import os
 from dataclasses import dataclass, replace
-from typing import Mapping, Optional, Tuple
+from typing import List, Mapping, Optional, Tuple
 
 from pants.base.build_environment import get_buildroot
 from pants.base.exception_sink import ExceptionSink
@@ -13,6 +13,7 @@
 from pants.base.specs_parser import SpecsParser
 from pants.base.workunit import WorkUnit
 from pants.build_graph.build_configuration import BuildConfiguration
+from pants.build_graph.lifecycle import SessionLifecycleHandler
 from pants.core.util_rules.pants_environment import PantsEnvironment
 from pants.engine.internals.native import Native
 from pants.engine.internals.scheduler import ExecutionError
@@ -55,6 +56,7 @@ class LocalPantsRunner:
     union_membership: UnionMembership
     profile_path: Optional[str]
     _run_tracker: RunTracker
+    _session_lifecycle_handlers: List[SessionLifecycleHandler]
 
     @classmethod
     def parse_options(
@@ -162,6 +164,16 @@ def create(
 
         profile_path = env.get("PANTS_PROFILE")
 
+        # Query registered extension lifecycle handlers to see if any wish to receive lifecycle events
+        # for this session.
+        session_lifecycle_handlers = []
+        for lifecycle_handler in build_config.lifecycle_handlers:
+            session_lifecycle_handler = lifecycle_handler.on_session_create(
+                build_root=build_root, options=options, specs=specs
+            )
+            if session_lifecycle_handler:
+                session_lifecycle_handlers.append(session_lifecycle_handler)
+
         return cls(
             build_root=build_root,
             options=options,
@@ -170,6 +182,7 @@ def create(
             graph_session=graph_session,
             union_membership=union_membership,
             profile_path=profile_path,
+            _session_lifecycle_handlers=session_lifecycle_handlers,
             _run_tracker=RunTracker.global_instance(),
         )
 
@@ -274,10 +287,19 @@ def run(self, start_time: float) -> ExitCode:
                 return help_printer.print_help()
 
             with streaming_reporter.session():
+                # Invoke the session lifecycle handlers, if any.
+                for session_lifecycle_handler in self._session_lifecycle_handlers:
+                    session_lifecycle_handler.on_session_start()
+
                 engine_result = PANTS_FAILED_EXIT_CODE
                 try:
                     engine_result = self._run_v2()
                 except Exception as e:
                     ExceptionSink.log_exception(e)
                 run_tracker_result = self._finish_run(engine_result)
+
+                # Invoke the session lifecycle handlers, if any.
+                for session_lifecycle_handler in self._session_lifecycle_handlers:
+                    session_lifecycle_handler.on_session_end(engine_result=engine_result)
+
             return self._merge_exit_codes(engine_result, run_tracker_result)

src/python/pants/build_graph/build_configuration.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, Type, Union
 
 from pants.build_graph.build_file_aliases import BuildFileAliases
+from pants.build_graph.lifecycle import ExtensionLifecycleHandler
 from pants.engine.rules import Rule, RuleIndex
 from pants.engine.target import Target
 from pants.engine.unions import UnionRule
@@ -26,6 +27,7 @@ class BuildConfiguration:
     rules: FrozenOrderedSet[Rule]
     union_rules: FrozenOrderedSet[UnionRule]
     target_types: FrozenOrderedSet[Type[Target]]
+    lifecycle_handlers: FrozenOrderedSet[ExtensionLifecycleHandler]
 
     @dataclass
     class Builder:
@@ -35,6 +37,9 @@ class Builder:
         _rules: OrderedSet = field(default_factory=OrderedSet)
         _union_rules: OrderedSet = field(default_factory=OrderedSet)
         _target_types: OrderedSet[Type[Target]] = field(default_factory=OrderedSet)
+        _lifecycle_handlers: OrderedSet[ExtensionLifecycleHandler] = field(
+            default_factory=OrderedSet
+        )
 
         def registered_aliases(self) -> BuildFileAliases:
             """Return the registered aliases exposed in BUILD files.
@@ -161,6 +166,25 @@ def register_target_types(
                 )
             self._target_types.update(target_types)
 
+        def register_lifecycle_handlers(
+            self, handlers: Union[typing.Iterable[ExtensionLifecycleHandler], typing.Any]
+        ):
+            if not isinstance(handlers, Iterable):
+                raise TypeError(
+                    f"The entrypoint `lifecycle_handlers` must return an iterable. Given {repr(handlers)}"
+                )
+            bad_elements = [
+                handler
+                for handler in handlers
+                if not isinstance(handler, ExtensionLifecycleHandler)
+            ]
+            if bad_elements:
+                raise TypeError(
+                    "Every element of the entrypoint `lifecycle_handlers` must be a subclass of "
+                    f"{ExtensionLifecycleHandler.__name__}. Bad elements: {bad_elements}."
+                )
+            self._lifecycle_handlers.update(handlers)
+
         def create(self) -> "BuildConfiguration":
             registered_aliases = BuildFileAliases(
                 objects=self._exposed_object_by_alias.copy(),
@@ -172,4 +196,5 @@ def create(self) -> "BuildConfiguration":
                 rules=FrozenOrderedSet(self._rules),
                 union_rules=FrozenOrderedSet(self._union_rules),
                 target_types=FrozenOrderedSet(self._target_types),
+                lifecycle_handlers=FrozenOrderedSet(self._lifecycle_handlers),
             )

src/python/pants/build_graph/build_configuration_test.py

@@ -2,10 +2,14 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import unittest
+from typing import Optional
 
+from pants.base.specs import Specs
 from pants.build_graph.build_configuration import BuildConfiguration
 from pants.build_graph.build_file_aliases import BuildFileAliases
+from pants.build_graph.lifecycle import ExtensionLifecycleHandler, SessionLifecycleHandler
 from pants.engine.unions import UnionRule, union
+from pants.option.options import Options
 from pants.util.frozendict import FrozenDict
 from pants.util.ordered_set import FrozenOrderedSet
 
@@ -64,3 +68,31 @@ class B:
         self.bc_builder.register_rules([union_a])
         self.bc_builder.register_rules([union_b])
         assert self.bc_builder.create().union_rules == FrozenOrderedSet([union_a, union_b])
+
+    def test_register_lifecycle_handlers(self) -> None:
+        class FooHandler(ExtensionLifecycleHandler):
+            def on_session_create(
+                self,
+                *,
+                build_root: str,
+                options: Options,
+                specs: Specs,
+            ) -> Optional[SessionLifecycleHandler]:
+                return None
+
+        class BarHandler(ExtensionLifecycleHandler):
+            def on_session_create(
+                self,
+                *,
+                build_root: str,
+                options: Options,
+                specs: Specs,
+            ) -> Optional[SessionLifecycleHandler]:
+                return None
+
+        self.bc_builder.register_lifecycle_handlers([FooHandler()])
+        self.bc_builder.register_lifecycle_handlers([BarHandler()])
+        bc = self.bc_builder.create()
+        handlers = list(bc.lifecycle_handlers)
+        assert isinstance(handlers[0], FooHandler)
+        assert isinstance(handlers[1], BarHandler)

src/python/pants/build_graph/lifecycle.py

@@ -0,0 +1,23 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from typing import Optional
+
+from pants.base.exiter import ExitCode
+from pants.base.specs import Specs
+from pants.option.options import Options
+
+
+class SessionLifecycleHandler:
+    def on_session_start(self):
+        pass
+
+    def on_session_end(self, *, engine_result: ExitCode):
+        pass
+
+
+class ExtensionLifecycleHandler:
+    # Returns a SessionLifecycleHandler that will receive lifecycle events for the session.
+    def on_session_create(
+        self, *, build_root: str, options: Options, specs: Specs
+    ) -> Optional[SessionLifecycleHandler]:
+        pass

src/python/pants/init/extension_loader.py

@@ -96,6 +96,9 @@ def load_plugins(
         if "rules" in entries:
             rules = entries["rules"].load()()
             build_configuration.register_rules(rules)
+        if "lifecycle_handlers" in entries:
+            lifecycle_handlers = entries["lifecycle_handlers"].load()()
+            build_configuration.register_lifecycle_handlers(lifecycle_handlers)
         loaded[dist.as_requirement().key] = dist
 
 
@@ -151,3 +154,6 @@ def invoke_entrypoint(name):
     rules = invoke_entrypoint("rules")
     if rules:
         build_configuration.register_rules(rules)
+    lifecycle_handlers = invoke_entrypoint("lifecycle_handlers")
+    if lifecycle_handlers:
+        build_configuration.register_lifecycle_handlers(lifecycle_handlers)