perf: Remove python coroutine/trampoline overhead in awaits for ~22% faster dependencies goal

src/python/pants/engine/internals/native_engine.pyi

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Coroutine, Generator, Iterable, Iterator, Mapping, Sequence
 from datetime import datetime
 from enum import Enum
 from io import RawIOBase
@@ -1155,12 +1155,14 @@ def hash_prefix_zero_bits(item: str) -> int: ...
 # ------------------------------------------------------------------------------
 
 _Output = TypeVar("_Output")
+_Output_co = TypeVar("_Output_co", covariant=True)
 _Input = TypeVar("_Input")
 
-class PyGeneratorResponseCall:
+class Call(Generic[_Output_co]):
     rule_id: str
     output_type: type
-    inputs: Sequence[Any]
+    args: tuple[Any, ...]
+    implicit_args: dict[Any, type]
 
     @overload
     def __init__(
@@ -1192,6 +1194,40 @@ class PyGeneratorResponseCall:
         input_arg0: type[_Input] | _Input,
         input_arg1: _Input | None = None,
     ) -> None: ...
+    def __await__(self) -> Generator[Any, None, _Output_co]: ...
+    def __repr__(self) -> str: ...
+
+class _Concurrently(Generic[_Output_co]):
+    def __init__(
+        self, calls: tuple[Coroutine[Any, Any, Any] | Call[Any] | _Concurrently[Any], ...]
+    ) -> None: ...
+    def __await__(self) -> Generator[_Concurrently[_Output_co], None, _Output_co]: ...
+    @property
+    def calls(
+        self,
+    ) -> tuple[Coroutine[Any, Any, Any] | Call[Any] | _Concurrently[Any], ...]: ...
+
+class RuleCallTrampoline(Generic[_Output]):
+    """The callable `@rule` returns. Captures `rule_id` and `output_type` at decoration time so
+    each invocation constructs the already-awaitable `Call` directly.
+    `__getattribute__` forwards `__doc__` and other introspection attrs to the wrapped function.
+    """
+
+    rule_id: str
+    output_type: type[_Output]
+    rule: Any
+    __wrapped__: Callable[..., Any]
+
+    def __init__(
+        self,
+        rule_id: str,
+        output_type: type[_Output],
+        wrapped: Callable[..., Any],
+        rule: Any,
+    ) -> None: ...
+    def __call__(
+        self, *args: Any, __implicitly: Sequence[Any] = (), **kwargs: Any
+    ) -> Call[_Output]: ...
 
 # ------------------------------------------------------------------------------
 # (uncategorized)

src/python/pants/engine/internals/rule_visitor.py

@@ -16,6 +16,7 @@
 import typing_extensions
 
 from pants.base.exceptions import RuleTypeError
+from pants.engine.internals.native_engine import RuleCallTrampoline
 from pants.engine.internals.selectors import (
     AwaitableConstraints,
     concurrently,
@@ -177,6 +178,11 @@ def _lookup_return_type(func: Callable, check: bool = False) -> Any:
 
 class _AwaitableCollector(ast.NodeVisitor):
     def __init__(self, func: Callable):
+        # `func` may be a RuleCallTrampoline (the return value of an `@rule`-decorated
+        # function). `inspect.getsource` and friends only know about real Python functions,
+        # so follow `__wrapped__` to reach the underlying implementation.
+        if isinstance(func, RuleCallTrampoline):
+            func = func.__wrapped__
         self.func = func
         source = inspect.getsource(func) or "<string>"
         beginning_indent = _get_starting_indent(source)
@@ -314,9 +320,9 @@ def _get_byname_awaitable(
     def visit_Call(self, call_node: ast.Call) -> None:
         func = self._lookup(call_node.func)
         if func is not None:
-            if (inspect.isfunction(func) or isinstance(func, RuleDescriptor)) and (
-                rule_id := getattr(func, "rule_id", None)
-            ) is not None:
+            if (
+                inspect.isfunction(func) or isinstance(func, (RuleDescriptor, RuleCallTrampoline))
+            ) and (rule_id := getattr(func, "rule_id", None)) is not None:
                 # Is a direct `@rule` call.
                 self.awaitables.append(self._get_byname_awaitable(rule_id, func, call_node))
             elif inspect.iscoroutinefunction(func):

src/python/pants/engine/internals/scheduler_test.py

@@ -474,12 +474,6 @@ def test_trace_includes_nested_exception_traceback() -> None:
           File LOCATION-INFO, in catch_and_reraise
             return await raise_an_exception(SomeInput(outer_input.s))
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-          File LOCATION-INFO, in wrapper
-            return await call
-                   ^^^^^^^^^^
-          File LOCATION-INFO, in __await__
-            result = yield self
-                     ^^^^^^^^^^
           File LOCATION-INFO, in raise_an_exception
             raise Exception(some_input.s)
         Exception: asdf