perf(rules): cut custom-rule host overhead

src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py

@@ -15,6 +15,7 @@
 
 from sqlbuild.compiler.compile.models import CompiledProject
 from sqlbuild.compiler.fact_cache.main.code_identity import compiled_code_identity
+from sqlbuild.compiler.scopes.models import ScopeIndex
 from sqlbuild.rule_engine._helpers.engine.custom_rule_evidence import (
     custom_rule_implementation_fingerprint,
     custom_rule_import_closure,
@@ -81,9 +82,10 @@ def evaluate_custom_rules_cached(
     custom: tuple[Rule, ...] = tuple(sorted(rules, key=lambda rule: rule.code))
     if not custom:
         return CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0)
+    facts_project: CompiledProject = host_project(project)
     if not config.cache.enabled:
         uncached: CustomHostRun = run_custom_hosts(
-            project=_host_project(project),
+            project=facts_project,
             config=config,
             project_dir=project_dir,
             dialect=dialect,
@@ -104,7 +106,7 @@ def evaluate_custom_rules_cached(
             custom_ms=_elapsed_ms(started),
         )
     digests: FactDigests = FactDigests(
-        views=_views_factory(project=project, project_dir=project_dir, dialect=dialect)
+        views=_views_factory(project=facts_project, project_dir=project_dir, dialect=dialect)
     )
     subjects: dict[str, tuple[_Subject, ...]] = _subjects(project=project, rules=custom)
     identities: dict[str, str] = {
@@ -150,7 +152,7 @@ def evaluate_custom_rules_cached(
     )
     if plan:
         run: CustomHostRun = run_custom_hosts(
-            project=_host_project(project),
+            project=facts_project,
             config=config,
             project_dir=project_dir,
             dialect=dialect,
@@ -187,9 +189,7 @@ def _views_factory(
     *, project: CompiledProject, project_dir: Path, dialect: str
 ) -> Callable[[], RuleFactViews]:
     def build() -> RuleFactViews:
-        return build_rule_fact_views(
-            project=_host_project(project), project_dir=project_dir, dialect=dialect
-        )
+        return build_rule_fact_views(project=project, project_dir=project_dir, dialect=dialect)
 
     return build
 
@@ -268,7 +268,9 @@ def _readset_digest(
     return memo[index]
 
 
-def _host_project(project: CompiledProject) -> CompiledProject:
+def host_project(project: CompiledProject) -> CompiledProject:
+    """Drop project state no Rule fact reads, before digests and host payloads use it."""
+
     return replace(
         project,
         binding_catalog=None,
@@ -278,6 +280,11 @@ def _host_project(project: CompiledProject) -> CompiledProject:
         hook_functions=(),
         materialization_files=(),
         external_sql_reference_resolver=None,
+        sql_scenarios=(),
+        sql_hook_files=(),
+        diagnostics=(),
+        unmatched_literal_sql_relations=(),
+        scope_index=ScopeIndex(),
     )
 
 

src/sqlbuild/rule_engine/_helpers/engine/fact_replay.py

@@ -9,6 +9,7 @@
 
 import orjson
 
+from sqlbuild.rule_engine.classes.shared_fact_encoder import SharedFactEncoder
 from sqlbuild.rule_engine.classes.sql_document import SqlDocument
 from sqlbuild.rule_engine.constants import (
     FACT_AUDITS_ALL,
@@ -98,6 +99,7 @@
     FACT_TREE_READ_TEXT: 1,
 }
 _TREE_FACTS: frozenset[str] = frozenset({FACT_TREE_PATHS, *_TEXT_FACTS})
+_SHARED_FACTS: frozenset[str] = frozenset(_GLOBAL_FACTS) - _TREE_FACTS
 
 
 def fact_key_payload(key: tuple[object, ...]) -> list[str]:
@@ -123,18 +125,31 @@ def full_fact_keys(views: RuleFactViews) -> tuple[FactKey, ...]:
     return tuple(keys)
 
 
-def fact_outcome_digest(*, views: RuleFactViews, key: FactKey) -> str:
-    """Recompute one recorded fact and digest its value or raised error."""
+def fact_outcome_digest(
+    *, views: RuleFactViews, key: FactKey, shared: SharedFactEncoder | None = None
+) -> str:
+    """Recompute one recorded fact and digest its value, sharing whole-project encodings."""
 
     try:
         value: object = _fact_value(views=views, key=key)
     except FactDigestError:
         raise
     except Exception as error:
         return fact_error_digest(error)
+    if shared is not None and key[0] in _SHARED_FACTS:
+        try:
+            return hashlib.sha256(shared.encode(("value", value))).hexdigest()
+        except (TypeError, FactDigestError) as error:
+            raise FactDigestError(str(error)) from error
     return fact_value_digest(value)
 
 
+def shared_fact_encoder() -> SharedFactEncoder:
+    """Return an encoder for whole-project facts that share compiler objects."""
+
+    return SharedFactEncoder(options=_ORJSON_OPTIONS, encode_default=_encode_default)
+
+
 def fact_value_digest(value: object) -> str:
     """Digest one fact value exactly as a recomputed fact would digest it."""
 

src/sqlbuild/rule_engine/_helpers/host/custom_host.py

@@ -3,14 +3,16 @@
 from __future__ import annotations
 
 import contextlib
+import gc
 import io
 import json
 import os
 import pickle
 import sys
+import time
 from dataclasses import asdict
 from pathlib import Path
-from typing import Any
+from typing import Any, BinaryIO
 
 from sqlbuild.compiler.compile.models import CompiledProject
 from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue
@@ -20,8 +22,11 @@
 from sqlbuild.rule_engine.constants import (
     CUSTOM_HOST_CANCELLED_EXIT_CODE,
     CUSTOM_HOST_CANCELLED_MESSAGE,
+    CUSTOM_HOST_EXEC_OS_NAME,
     CUSTOM_HOST_HASH_SEED,
+    CUSTOM_HOST_INPUT_POLL_SECONDS,
     CUSTOM_HOST_INPUT_TUPLE_SIZE,
+    CUSTOM_HOST_INPUT_WAIT_SECONDS,
     CUSTOM_HOST_LAUNCH_MODULE,
     CUSTOM_HOST_MAX_TRACKED_READS,
     CUSTOM_HOST_PROTOCOL_VERSION,
@@ -61,7 +66,7 @@ def main() -> int:
         cancel_marker: Path | None = _cancel_marker(payload)
         if cancel_marker is not None and cancel_marker.exists():
             raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE)
-        project, config = _decode_inputs(payload)
+        project, config = _decode_inputs(payload=payload, cancel_marker=cancel_marker)
         project_dir: Path = Path(str(payload["project_dir"])).resolve()
         dialect: str = str(payload.get("dialect", "generic"))
         verify_determinism: bool = payload.get("verify_determinism") is True
@@ -111,18 +116,27 @@ def main() -> int:
     return 0
 
 
-def _decode_inputs(payload: dict[str, Any]) -> tuple[CompiledProject, RulesConfig]:
+def _decode_inputs(
+    *, payload: dict[str, Any], cancel_marker: Path | None
+) -> tuple[CompiledProject, RulesConfig]:
     project_dir: Path = Path(str(payload["project_dir"])).resolve()
     input_root: Path = (project_dir / "target" / "rules-cache" / "host-inputs").resolve()
     input_path: Path = Path(str(payload["project_pickle_path"]))
-    if (
-        not input_path.is_absolute()
-        or input_path.is_symlink()
-        or not input_path.resolve().is_relative_to(input_root)
-    ):
+    if not input_path.is_absolute() or not input_path.parent.resolve().is_relative_to(input_root):
+        raise RulesError("custom host project payload path is invalid")
+    await_inputs(
+        path=input_path,
+        cancel_markers=tuple(
+            marker
+            for marker in (cancel_marker, _input_marker(payload=payload, key="abandoned_marker"))
+            if marker is not None
+        ),
+        parent_pid=_expected_parent(payload),
+    )
+    if input_path.is_symlink() or not input_path.resolve().is_relative_to(input_root):
         raise RulesError("custom host project payload path is invalid")
     with input_path.open("rb") as handle:
-        decoded: object = pickle.load(handle)
+        decoded: object = _load_long_lived(handle)
     if (
         not isinstance(decoded, tuple)
         or len(decoded) != CUSTOM_HOST_INPUT_TUPLE_SIZE
@@ -133,20 +147,62 @@ def _decode_inputs(payload: dict[str, Any]) -> tuple[CompiledProject, RulesConfi
     return decoded
 
 
-def _cancel_marker(payload: dict[str, Any]) -> Path | None:
-    """Return the run's cancellation marker, which must live beside the project payload."""
+def await_inputs(*, path: Path, cancel_markers: tuple[Path, ...], parent_pid: int) -> None:
+    """Wait for the payload the parent publishes after starting this host, while it still runs."""
+
+    deadline: float = time.monotonic() + CUSTOM_HOST_INPUT_WAIT_SECONDS
+    while not path.exists():
+        if any(marker.exists() for marker in cancel_markers):
+            raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE)
+        if os.getppid() != parent_pid:
+            raise RulesError("custom host parent exited before writing the project payload")
+        if time.monotonic() > deadline:
+            raise RulesError("custom host project payload was never written")
+        time.sleep(CUSTOM_HOST_INPUT_POLL_SECONDS)
+
+
+def _expected_parent(payload: dict[str, Any]) -> int:
+    """Return the parent that started this host; a relaying launcher stands in for it."""
+
+    value: object = payload.get("parent_pid")
+    if os.name == CUSTOM_HOST_EXEC_OS_NAME and isinstance(value, int) and value > 0:
+        return value
+    return os.getppid()
+
 
-    value: object = payload.get("cancel_marker")
+def _input_marker(*, payload: dict[str, Any], key: str) -> Path | None:
+    """Return a signal file the run placed beside the project payload, if it named one."""
+
+    value: object = payload.get(key)
     if value is None:
         return None
     project_dir: Path = Path(str(payload["project_dir"])).resolve()
     input_root: Path = (project_dir / "target" / "rules-cache" / "host-inputs").resolve()
     marker: Path = Path(str(value))
-    if not marker.is_absolute() or not marker.resolve().is_relative_to(input_root):
-        raise RulesError("custom host cancellation marker path is invalid")
+    if not marker.is_absolute() or not marker.parent.resolve().is_relative_to(input_root):
+        raise RulesError(f"custom host {key} path is invalid")
     return marker
 
 
+def _load_long_lived(handle: BinaryIO) -> Any:
+    """Unpickle the host-lifetime project without cyclic collection, then freeze it."""
+
+    was_enabled: bool = gc.isenabled()
+    gc.disable()
+    try:
+        return pickle.load(handle)
+    finally:
+        gc.freeze()
+        if was_enabled:
+            gc.enable()
+
+
+def _cancel_marker(payload: dict[str, Any]) -> Path | None:
+    """Return the run's cancellation marker, which must live beside the project payload."""
+
+    return _input_marker(payload=payload, key="cancel_marker")
+
+
 def _decode_plan(value: object) -> CustomRulePlan:
     if not isinstance(value, dict):
         raise RulesError("custom host plan must be an object")