Support catching @rule errors (#17911)

* reviving #10954, superseding #17745, closing #17910 

Co-authored-by: Danny McClanahan <1305167+cosmicexplorer@users.noreply.github.com>

src/python/pants/backend/python/goals/setup_py_test.py

@@ -55,6 +55,7 @@
 )
 from pants.backend.python.util_rules import dists, python_sources
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.base.exceptions import IntrinsicError
 from pants.core.goals.package import BuiltPackage
 from pants.core.target_types import FileTarget, ResourcesGeneratorTarget, ResourceTarget
 from pants.core.target_types import rules as core_target_types_rules
@@ -615,7 +616,7 @@ def test_generate_long_description_field_from_non_existing_file(
     assert_chroot_error(
         chroot_rule_runner,
         Address("src/python/foo", target_name="foo-dist"),
-        Exception,
+        IntrinsicError,
     )
 
 

src/python/pants/base/exceptions.py

@@ -5,6 +5,12 @@
 
 from typing import TYPE_CHECKING
 
+from pants.engine.internals.native_engine import EngineError as EngineError  # noqa: F401
+from pants.engine.internals.native_engine import (  # noqa: F401
+    IncorrectProductError as IncorrectProductError,
+)
+from pants.engine.internals.native_engine import IntrinsicError as IntrinsicError  # noqa: F401
+
 if TYPE_CHECKING:
     from pants.engine.internals.native_engine import PyFailure
 
@@ -38,6 +44,13 @@ class MappingError(Exception):
 class NativeEngineFailure(Exception):
     """A wrapper around a `Failure` instance.
 
+    The failure instance being wrapped can come from an exception raised in a rule. When this
+    failure is returned to a requesting rule it is first unwrapped so the original exception will be
+    presented in the rule, thus the `NativeEngineFailure` exception will not be seen in rule code.
+
+    This is different from the other `EngineError` based exceptions which doesn't originate from
+    rule code.
+
     TODO: This type is defined in Python because pyo3 doesn't support declaring Exceptions with
     additional fields. See https://github.com/PyO3/pyo3/issues/295
     """

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

@@ -11,6 +11,7 @@
 import pytest
 
 from pants.backend.python.target_types import PythonSourcesGeneratorTarget
+from pants.base.exceptions import IntrinsicError
 from pants.base.specs import Specs
 from pants.base.specs_parser import SpecsParser
 from pants.engine.engine_aware import EngineAwareParameter, EngineAwareReturnType
@@ -21,6 +22,7 @@
     Digest,
     DigestContents,
     FileContent,
+    MergeDigests,
     Snapshot,
 )
 from pants.engine.internals.engine_testutil import (
@@ -40,8 +42,9 @@
 from pants.engine.unions import UnionRule, union
 from pants.goal.run_tracker import RunTracker
 from pants.testutil.option_util import create_options_bootstrapper
-from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.testutil.rule_runner import QueryRule, RuleRunner, engine_error
 from pants.util.logging import LogLevel
+from pants.util.strutil import softwrap
 
 
 class A:
@@ -1006,3 +1009,44 @@ async def for_member() -> str:
     )
 
     assert "yep" == rule_runner.request(str, [])
+
+
+@dataclass(frozen=True)
+class FileInput:
+    filename: str
+
+
+@dataclass(frozen=True)
+class MergedOutput:
+    digest: Digest
+
+
+class MergeErr(Exception):
+    pass
+
+
+@rule
+async def catch_merge_digests_error(file_input: FileInput) -> MergedOutput:
+    # Create two separate digests writing different contents to the same file path.
+    input_1 = CreateDigest((FileContent(path=file_input.filename, content=b"yes"),))
+    input_2 = CreateDigest((FileContent(path=file_input.filename, content=b"no"),))
+    digests = await MultiGet(Get(Digest, CreateDigest, input_1), Get(Digest, CreateDigest, input_2))
+    try:
+        merged = await Get(Digest, MergeDigests(digests))
+    except IntrinsicError as e:
+        raise MergeErr(f"error merging digests for input {file_input}: {e}")
+    return MergedOutput(merged)
+
+
+def test_catch_intrinsic_error() -> None:
+    rule_runner = RuleRunner(
+        rules=[catch_merge_digests_error, QueryRule(MergedOutput, (FileInput,))]
+    )
+    msg = softwrap(
+        """\
+        error merging digests for input FileInput(filename='some-file.txt'): Can only merge
+        Directories with no duplicates, but found 2 duplicate entries in :
+        """
+    )
+    with engine_error(MergeErr, contains=msg):
+        rule_runner.request(MergedOutput, (FileInput("some-file.txt"),))

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

@@ -20,7 +20,8 @@ from pants.engine.process import InteractiveProcess, InteractiveProcessResult
 # (core)
 # ------------------------------------------------------------------------------
 
-class PyFailure: ...
+class PyFailure:
+    def get_error(self) -> Exception | None: ...
 
 # ------------------------------------------------------------------------------
 # Address (parsing)
@@ -478,3 +479,14 @@ class PyThreadLocals:
 
 class PollTimeout(Exception):
     pass
+
+# Prefer to import these exception types from `pants.base.exceptions`
+
+class EngineError(Exception):
+    """Base exception used for errors originating from the native engine."""
+
+class IntrinsicError(EngineError):
+    """Exceptions raised for failures within intrinsic methods implemented in Rust."""
+
+class IncorrectProductError(EngineError):
+    """Exceptions raised when a rule's return value doesn't match its declared type."""

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

@@ -9,11 +9,12 @@
 
 import pytest
 
+from pants.base.exceptions import IncorrectProductError
 from pants.engine.internals.engine_testutil import remove_locations_from_traceback
 from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.rules import Get, rule
 from pants.engine.unions import UnionRule, union
-from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.testutil.rule_runner import QueryRule, RuleRunner, engine_error
 
 # -----------------------------------------------------------------------------------------------
 # Test params
@@ -248,6 +249,115 @@ def test_outlined_get() -> None:
     ) in str(exc.value.args[0])
 
 
+@dataclass(frozen=True)
+class SomeInput:
+    s: str
+
+
+@dataclass(frozen=True)
+class SomeOutput:
+    s: str
+
+
+@rule
+def raise_an_exception(some_input: SomeInput) -> SomeOutput:
+    raise Exception(some_input.s)
+
+
+@dataclass(frozen=True)
+class OuterInput:
+    s: str
+
+
+@rule
+async def catch_an_exception(outer_input: OuterInput) -> SomeOutput:
+    try:
+        return await Get(SomeOutput, SomeInput(outer_input.s))
+    except Exception as e:
+        return SomeOutput(str(e))
+
+
+@rule(desc="error chain")
+async def catch_and_reraise(outer_input: OuterInput) -> SomeOutput:
+    """This rule is used in a dedicated test only, so does not conflict with
+    `catch_an_exception`."""
+    try:
+        return await Get(SomeOutput, SomeInput(outer_input.s))
+    except Exception as e:
+        raise Exception("nested exception!") from e
+
+
+class InputWithNothing:
+    pass
+
+
+GLOBAL_FLAG: bool = True
+
+
+@rule
+def raise_an_exception_upon_global_state(input_with_nothing: InputWithNothing) -> SomeOutput:
+    if GLOBAL_FLAG:
+        raise Exception("global flag is set!")
+    return SomeOutput("asdf")
+
+
+@rule
+def return_a_wrong_product_type(input_with_nothing: InputWithNothing) -> A:
+    return B()  # type: ignore[return-value]
+
+
+@rule
+async def catch_a_wrong_product_type(input_with_nothing: InputWithNothing) -> B:
+    try:
+        _ = await Get(A, InputWithNothing, input_with_nothing)
+    except IncorrectProductError as e:
+        raise Exception(f"caught product type error: {e}")
+    return B()
+
+
+@pytest.fixture
+def rule_error_runner() -> RuleRunner:
+    return RuleRunner(
+        rules=(
+            consumes_a_and_b,
+            QueryRule(str, (A, B)),
+            transitive_b_c,
+            QueryRule(str, (A, C)),
+            transitive_coroutine_rule,
+            QueryRule(D, (C,)),
+            boolean_and_int,
+            QueryRule(A, (int, bool)),
+            raise_an_exception,
+            QueryRule(SomeOutput, (SomeInput,)),
+            catch_an_exception,
+            QueryRule(SomeOutput, (OuterInput,)),
+            raise_an_exception_upon_global_state,
+            QueryRule(SomeOutput, (InputWithNothing,)),
+            return_a_wrong_product_type,
+            QueryRule(A, (InputWithNothing,)),
+            catch_a_wrong_product_type,
+            QueryRule(B, (InputWithNothing,)),
+        )
+    )
+
+
+def test_catch_inner_exception(rule_error_runner: RuleRunner) -> None:
+    assert rule_error_runner.request(SomeOutput, [OuterInput("asdf")]) == SomeOutput("asdf")
+
+
+def test_exceptions_uncached(rule_error_runner: RuleRunner) -> None:
+    global GLOBAL_FLAG
+    with engine_error(Exception, contains="global flag is set!"):
+        rule_error_runner.request(SomeOutput, [InputWithNothing()])
+    GLOBAL_FLAG = False
+    assert rule_error_runner.request(SomeOutput, [InputWithNothing()]) == SomeOutput("asdf")
+
+
+def test_incorrect_product_type(rule_error_runner: RuleRunner) -> None:
+    with engine_error(Exception, contains="caught product type error"):
+        rule_error_runner.request(B, [InputWithNothing()])
+
+
 # -----------------------------------------------------------------------------------------------
 # Test tracebacks
 # -----------------------------------------------------------------------------------------------
@@ -264,11 +374,7 @@ def nested_raise() -> A:
 
 
 def test_trace_includes_rule_exception_traceback() -> None:
-    rule_runner = RuleRunner(rules=[nested_raise, QueryRule(A, [])])
-    with pytest.raises(ExecutionError) as exc:
-        rule_runner.request(A, [])
-    normalized_traceback = remove_locations_from_traceback(str(exc.value))
-    assert normalized_traceback == dedent(
+    normalized_traceback = dedent(
         f"""\
          1 Exception encountered:
 
@@ -288,6 +394,59 @@ def test_trace_includes_rule_exception_traceback() -> None:
          """
     )
 
+    rule_runner = RuleRunner(rules=[nested_raise, QueryRule(A, [])])
+    with engine_error(Exception, contains=normalized_traceback, normalize_tracebacks=True):
+        rule_runner.request(A, [])
+
+
+def test_trace_includes_nested_exception_traceback() -> None:
+    normalized_traceback = dedent(
+        f"""\
+        1 Exception encountered:
+
+        Engine traceback:
+          in select
+            ..
+          in {__name__}.{catch_and_reraise.__name__}
+            error chain
+          in {__name__}.{raise_an_exception.__name__}
+            ..
+
+        Traceback (most recent call last):
+          File LOCATION-INFO, in raise_an_exception
+            raise Exception(some_input.s)
+        Exception: asdf
+
+        During handling of the above exception, another exception occurred:
+
+        Traceback (most recent call last):
+          File LOCATION-INFO, in catch_and_reraise
+            return await Get(SomeOutput, SomeInput(outer_input.s))
+          File LOCATION-INFO, in __await__
+            result = yield self
+        Exception: asdf
+
+        The above exception was the direct cause of the following exception:
+
+        Traceback (most recent call last):
+          File LOCATION-INFO, in native_engine_generator_send
+            res = rule.send(arg) if err is None else rule.throw(throw or err)
+          File LOCATION-INFO, in catch_and_reraise
+            raise Exception("nested exception!") from e
+        Exception: nested exception!
+        """
+    )
+
+    rule_runner = RuleRunner(
+        rules=[
+            raise_an_exception,
+            catch_and_reraise,
+            QueryRule(SomeOutput, (OuterInput,)),
+        ]
+    )
+    with engine_error(Exception, contains=normalized_traceback, normalize_tracebacks=True):
+        rule_runner.request(SomeOutput, [OuterInput("asdf")])
+
 
 # -----------------------------------------------------------------------------------------------
 # Test unhashable types