From 82971005ec9ceeeee6a55a3ef6de5f5793a47823 Mon Sep 17 00:00:00 2001 From: Micah Smith Date: Mon, 12 Oct 2026 17:52:05 -0400 Subject: [PATCH] refactor(autoland): merge land runs as a stack through a plan optimization policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Phase 4 of the autoland step-abstraction stack. After the plan is parsed, the engine runs a plan optimization pass, driven by a registry of policies, and executes the result. Native stack merging moves out of the engine into the first such policy, so the base steps are fully independent of each other and of the engine. - **Execution units** (new `autoland/optimize/` package). The engine executes a list of `ExecutionUnit`s covering the rest of the plan, from `current_step`. A unit covers one or more consecutive plan step indices (`indices: range`) and has `execute(uc: UnitContext) -> str` (the abort reason, or "") and `preview() -> list[str]`. The default `StepUnit` wraps one step and just executes it. - **`UnitContext`** is what a unit runs its steps with: `run`, `opts`, `common`, `services`, and three engine callbacks: - `enter(index)`: make the step current and checkpoint, exactly as the engine loop did before each step. - `execute_step(index)`: build the step's `StepContext`, execute it, and record the landed SHA it reports. - `checkpoint()`: save the run. So a composite unit still checkpoints and executes each of its steps exactly as a `StepUnit` would. - **Optimization policies.** An `OptimizationPolicy` is a `Protocol` with a `name`, `applies(opts) -> bool`, and `optimize(units, run, opts) -> units`. `@register_policy` adds a policy class to the ordered `POLICIES` registry. A duplicate name is an error. `optimize_plan(run, opts, policies=None)` starts from one `StepUnit` per remaining step, applies each applicable policy in registration order, and raises if a policy's result doesn't cover the remaining steps in order. To add a policy, write the class, decorate it, and import its module from `optimize/__init__.py`. `execute_plan` and `print_execution_preview` take an optional `policies` registry, which tests use to plug in their own. - **Policies hold the cross-step knowledge.** A policy may look at step types and at the run. The steps never know about each other or about policies. - **`NativeStackMergePolicy`** (`optimize/stack_merge.py`, on with `--merge-as-stack`) groups each run of two or more consecutive land steps whose PRs are still open into a `StackMergeUnit`. Any other step ends a run: a `w` or `c` step, an already-landed `l`, or a merged PR. A run of one is left alone. The unit performs the stack merge with `native_stack.land_as_native_stack` and then executes each member step itself, so per-step output, skips, and landed-SHA reporting are unchanged. If the merge falls through, the steps it covered land one at a time without retrying it, the same semantics as the old `unstacked_through`. - **Re-validation at execution time.** Before each of its steps (until a merge falls through), a `StackMergeUnit` recomputes, from the live stack, which of its following steps still land an open PR, and only stack-merges those if there are at least two. This matches today's behavior, where the run was evaluated when its first step was reached: a PR merged elsewhere in the meantime shortens the run. - **The optimized plan is never persisted.** It exists only for one execution. The state file still holds the user's plan, byte-identical, and `current_step` is still an index into it. Every fresh, resumed, or replanned run optimizes again from `current_step`. - **The preview derives from the units.** `print_native_stack_runs` is replaced by `engine.print_execution_preview`, which prints each unit's `preview()` lines from the same `optimize_plan` that `execute_plan` uses, so the preview and the execution can't disagree. The text is unchanged. `commands.py` calls it in the same three places. - The temporary `_NativeStackMerge.before_step` hook is gone from the engine. `native_stack.py` keeps GitHub's side of a stack merge (`land_as_native_stack` and its helpers). The run detection that peeked at neighboring steps (`native_stack_run`, `native_stack_runs`) moves into the policy. No user-visible change. Console output, checkpoints, and status output are byte-for-byte identical. ## Test plan - [x] The golden autoland scenarios (`tests/test_autoland_golden.py`, including `merge_as_stack` and `stack_merge_falls_through`) pass with no golden files regenerated. - [x] New `tests/test_autoland_optimize.py`: - Registry: policies run in registration order; registering a name twice is an error; a policy that drops a step is rejected. - Native-stack policy grouping: runs split by `w`/`c` steps (moved from `test_autoland.py`); already-landed steps and merged PRs left out of runs; runs of one left as `StepUnit`s; runs start at `current_step`; no grouping with `merge_as_stack` off; the preview text for each run. - `StackMergeUnit` execution: a stack merge followed by each step, with each step entered (checkpointed as current) before its work; a fall-through that lands the rest one at a time without retrying; a stack-merge abort stopping before the step; a step abort stopping the unit. - Re-validation: a PR merged after optimization shortens the merged run; a run shortened to one PR is not stack-merged. - Extension point: a test-only policy, registered from the test into its own registry, reshapes both the preview and the execution through `execute_plan`, with one checkpoint per step. - [x] The full suite passes with an empty `GH_CONFIG_DIR` and `GH_TOKEN` unset. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 5.5 (1M context) stack-info: PR: https://github.com/micahjsmith/stack-pr/pull/72, branch: micahjsmith/stack/6 --- src/stack_pr/autoland/__init__.py | 4 +- src/stack_pr/autoland/commands.py | 7 +- src/stack_pr/autoland/engine.py | 125 +++--- src/stack_pr/autoland/native_stack.py | 50 +-- src/stack_pr/autoland/optimize/__init__.py | 35 ++ src/stack_pr/autoland/optimize/base.py | 178 ++++++++ src/stack_pr/autoland/optimize/stack_merge.py | 157 +++++++ tests/test_autoland.py | 24 - tests/test_autoland_optimize.py | 411 ++++++++++++++++++ 9 files changed, 848 insertions(+), 143 deletions(-) create mode 100644 src/stack_pr/autoland/optimize/__init__.py create mode 100644 src/stack_pr/autoland/optimize/base.py create mode 100644 src/stack_pr/autoland/optimize/stack_merge.py create mode 100644 tests/test_autoland_optimize.py diff --git a/src/stack_pr/autoland/__init__.py b/src/stack_pr/autoland/__init__.py index 0f553c4..43e0dd3 100644 --- a/src/stack_pr/autoland/__init__.py +++ b/src/stack_pr/autoland/__init__.py @@ -32,7 +32,9 @@ - ``native_stack``: merging a run of land steps as one GitHub stack. - ``services``: the tools a step uses as it executes (console, GitHub, the waits, the repo), and their production bindings. -- ``engine``: executing a plan. +- ``optimize``: execution units, the plan optimization policies that group + steps into them (e.g. ``--merge-as-stack``), and the registry of policies. +- ``engine``: optimizing and executing a plan. - ``status``: the ``--status`` report. - ``commands``: the subparser, and the fresh / resume / replan flows. diff --git a/src/stack_pr/autoland/commands.py b/src/stack_pr/autoland/commands.py index 2c2ea7d..9df5021 100644 --- a/src/stack_pr/autoland/commands.py +++ b/src/stack_pr/autoland/commands.py @@ -15,7 +15,6 @@ from stack_pr.autoland import discovery, engine, runtime from stack_pr.autoland.display import _describe_step, _escape_markup, print_status from stack_pr.autoland.model import LandingContext, LandStep, PlanStep, StackEntry -from stack_pr.autoland.native_stack import print_native_stack_runs from stack_pr.autoland.options import AutolandOptions from stack_pr.autoland.plan import ( PLAN_SUFFIX, @@ -417,7 +416,7 @@ def _run_fresh(common: cli.CommonArgs, opts: AutolandOptions) -> None: ) print_status(ctx) - print_native_stack_runs(ctx, opts) + engine.print_execution_preview(ctx, opts) if opts.dry_run: runtime.console.print("\n[yellow]Dry run — exiting.[/yellow]") return @@ -496,7 +495,7 @@ def _run_resume(common: cli.CommonArgs, opts: AutolandOptions) -> None: return print_status(ctx) - print_native_stack_runs(ctx, opts) + engine.print_execution_preview(ctx, opts) runtime.console.print(f"[dim]State file: {sf_path}[/dim]\n") _install_signal_handler(ctx, checkpointer, worktree, opts) _finish( @@ -681,7 +680,7 @@ def _replan(common: cli.CommonArgs, opts: AutolandOptions, state_path: Path) -> ctx = LandingContext(stack=stack, plan=plan) print_status(ctx) - print_native_stack_runs(ctx, opts) + engine.print_execution_preview(ctx, opts) if lost: runtime.console.print( "\n[yellow]Done in the previous run, but not carried over (changed, " diff --git a/src/stack_pr/autoland/engine.py b/src/stack_pr/autoland/engine.py index 0d75bac..cb5f56b 100644 --- a/src/stack_pr/autoland/engine.py +++ b/src/stack_pr/autoland/engine.py @@ -1,9 +1,12 @@ -"""Executing a landing plan step by step. - -The engine owns the run: it walks the plan from ``ctx.current_step``, -checkpoints before each step, hands each step a ``StepContext``, records the -run-level facts a step reports in its ``StepResult``, and stops on an abort. -What a step does is up to the step (``AutolandStep.execute``). +"""Executing a landing plan. + +The engine owns the run. It optimizes the rest of the plan into execution +units (``optimize.optimize_plan``), then executes them in order. Through the +unit's ``UnitContext``, each step the unit runs is made the current step and +checkpointed before it starts, is handed a ``StepContext``, and has the +run-level facts it reports in its ``StepResult`` recorded. The run stops on +the first abort. What a step does is up to the step (``AutolandStep.execute``), +and what a unit does across its steps is up to the policy that built it. """ from __future__ import annotations @@ -11,14 +14,15 @@ import functools from stack_pr import cli +from stack_pr.autoland import runtime from stack_pr.autoland.display import _describe_step from stack_pr.autoland.model import LandingContext -from stack_pr.autoland.native_stack import land_as_native_stack, native_stack_run +from stack_pr.autoland.optimize import PolicyRegistry, UnitContext, optimize_plan from stack_pr.autoland.options import AutolandOptions from stack_pr.autoland.plan import landed_prefix_end from stack_pr.autoland.services import Services from stack_pr.autoland.state import AutolandCheckpointer -from stack_pr.autoland.steps import StepContext +from stack_pr.autoland.steps import StepContext, StepResult def execute_plan( @@ -28,12 +32,15 @@ def execute_plan( checkpointer: AutolandCheckpointer, *, services: Services | None = None, + policies: PolicyRegistry | None = None, ) -> bool: - """Execute the landing plan from ctx.current_step. Returns True on success.""" + """Execute the landing plan from ctx.current_step. Returns True on success. + + *policies* defaults to every registered optimization policy. + """ if services is None: services = Services.production() prefix_end = landed_prefix_end(ctx.plan) - native_stack = _NativeStackMerge(ctx, common, opts, checkpointer) def upcoming(index: int) -> list[str]: return [_describe_step(s, ctx) for s in ctx.plan[index + 1 :]] @@ -43,38 +50,63 @@ def landed_sha() -> str: _record_landed_sha(ctx, services.repo.landed_sha(common)) return ctx.last_landed_sha - for step_idx in range(ctx.current_step, len(ctx.plan)): - step = ctx.plan[step_idx] - ctx.current_step = step_idx + def checkpoint() -> None: checkpointer.save(ctx) - abort_reason = native_stack.before_step(step_idx) - if abort_reason: - return _abort(ctx, checkpointer, abort_reason) + def enter(index: int) -> None: + ctx.current_step = index + checkpoint() - result = step.execute( + def execute_step(index: int) -> StepResult: + result = ctx.plan[index].execute( StepContext( - index=step_idx, + index=index, total=len(ctx.plan), run=ctx, opts=opts, common=common, services=services, - checkpoint=lambda: checkpointer.save(ctx), - upcoming=functools.partial(upcoming, step_idx), + checkpoint=checkpoint, + upcoming=functools.partial(upcoming, index), landed_sha=landed_sha, - ends_landed_prefix=step_idx == prefix_end, + ends_landed_prefix=index == prefix_end, ) ) _record_landed_sha(ctx, result.landed_sha) - if result.aborted: - return _abort(ctx, checkpointer, result.abort_reason) + return result + + uc = UnitContext( + run=ctx, + opts=opts, + common=common, + services=services, + enter=enter, + execute_step=execute_step, + checkpoint=checkpoint, + ) + for unit in optimize_plan(ctx, opts, policies): + abort_reason = unit.execute(uc) + if abort_reason: + return _abort(ctx, checkpointer, abort_reason) ctx.current_step = len(ctx.plan) checkpointer.save(ctx) return True +def print_execution_preview( + ctx: LandingContext, + opts: AutolandOptions, + *, + policies: PolicyRegistry | None = None, +) -> None: + """Announce what the optimized plan does beyond executing its steps one by + one. Derived from the same units ``execute_plan`` executes.""" + for unit in optimize_plan(ctx, opts, policies): + for line in unit.preview(): + runtime.console.print(line) + + def _record_landed_sha(ctx: LandingContext, sha: str) -> None: # An unknown SHA ("") keeps the last one known. if sha: @@ -88,50 +120,3 @@ def _abort( ctx.aborted = True checkpointer.save(ctx) return False - - -class _NativeStackMerge: - """TEMPORARY: merge a run of consecutive land steps as one GitHub stack. - - With ``--merge-as-stack``, before each step the engine asks whether a run - of two or more open land steps starts there, and if so lands the run with - one stack merge (``native_stack.land_as_native_stack``). The land steps - then find their PRs merged and only skip. A run whose stack merge fell - through is landed one step at a time instead, without retrying the merge. - - This lives in the engine, not in ``LandStep``, because steps are - independent of each other. It is a stopgap: a plan-optimization pass that - groups land steps into one execution unit replaces it. - """ - - def __init__( - self, - ctx: LandingContext, - common: cli.CommonArgs, - opts: AutolandOptions, - checkpointer: AutolandCheckpointer, - ) -> None: - self._ctx = ctx - self._common = common - self._opts = opts - self._checkpointer = checkpointer - # The last step of a run whose stack merge fell through. - self._unstacked_through = -1 - - def before_step(self, step_idx: int) -> str: - """Stack-merge the run starting at *step_idx*, if any. Returns the - reason the plan must stop, or "".""" - if not self._opts.merge_as_stack or step_idx <= self._unstacked_through: - return "" - run_entries = native_stack_run(self._ctx, step_idx) - if not run_entries: - return "" - result = land_as_native_stack( - run_entries, ctx=self._ctx, common=self._common, opts=self._opts - ) - if result.abort_reason: - return result.abort_reason - if not result.landed: - self._unstacked_through = step_idx + len(run_entries) - 1 - self._checkpointer.save(self._ctx) - return "" diff --git a/src/stack_pr/autoland/native_stack.py b/src/stack_pr/autoland/native_stack.py index e50be56..8a02dd2 100644 --- a/src/stack_pr/autoland/native_stack.py +++ b/src/stack_pr/autoland/native_stack.py @@ -1,4 +1,8 @@ -"""Merging a run of consecutive land steps as one GitHub stack.""" +"""Merging a run of consecutive land steps as one GitHub stack. + +This is GitHub's side of a stack merge. Which steps one covers is decided by +the ``optimize.stack_merge`` policy. +""" from __future__ import annotations @@ -7,7 +11,7 @@ from stack_pr import cli from stack_pr.autoland import gh, runtime, waits -from stack_pr.autoland.model import LandingContext, LandStep, PRState, StackEntry +from stack_pr.autoland.model import LandingContext, PRState, StackEntry from stack_pr.autoland.options import AutolandOptions from stack_pr.autoland.waits import _poll_read @@ -16,48 +20,6 @@ # --------------------------------------------------------------------------- -def native_stack_run(ctx: LandingContext, start: int) -> list[StackEntry]: - """The PRs a stack merge starting at plan step *start* would land. - - That is the run of consecutive ``l`` steps from *start* whose PRs are still - open. A ``w`` or ``c`` step ends the run: it has to see the PRs above it - land separately from the ones below it. A run of fewer than two PRs gains - nothing from a stack merge, so it comes back empty. - """ - entries: list[StackEntry] = [] - for step in ctx.plan[start:]: - if not isinstance(step, LandStep) or step.already_landed: - break - entry = ctx.stack[step.entry_index] - if entry.state == PRState.MERGED: - break - entries.append(entry) - return entries if len(entries) > 1 else [] - - -def native_stack_runs(ctx: LandingContext) -> list[tuple[int, list[StackEntry]]]: - """Every ``(first step, PRs)`` run the rest of the plan would stack-merge.""" - runs = [] - index = ctx.current_step - while index < len(ctx.plan): - entries = native_stack_run(ctx, index) - if entries: - runs.append((index, entries)) - index += max(len(entries), 1) - return runs - - -def print_native_stack_runs(ctx: LandingContext, opts: AutolandOptions) -> None: - if not opts.merge_as_stack: - return - for first, entries in native_stack_runs(ctx): - prs = ", ".join(f"#{e.pr_number}" for e in entries) - runtime.console.print( - f"[dim]Steps {first + 1}-{first + len(entries)} ({prs}) will merge " - "together as a GitHub stack.[/dim]" - ) - - def _open_native_stack_prs(native_stack: dict) -> list[int]: """The numbers of a GitHub stack's unmerged PRs, bottom first.""" return [ diff --git a/src/stack_pr/autoland/optimize/__init__.py b/src/stack_pr/autoland/optimize/__init__.py new file mode 100644 index 0000000..5512ff7 --- /dev/null +++ b/src/stack_pr/autoland/optimize/__init__.py @@ -0,0 +1,35 @@ +"""Plan optimization: execution units and the policies that build them. + +Importing this package imports every policy module, so each policy is in the +registry before a plan is executed. +""" + +from __future__ import annotations + +from stack_pr.autoland.optimize.base import ( + POLICIES, + ExecutionUnit, + OptimizationPolicy, + PolicyRegistry, + StepUnit, + UnitContext, + optimize_plan, + register_policy, +) +from stack_pr.autoland.optimize.stack_merge import ( + NativeStackMergePolicy, + StackMergeUnit, +) + +__all__ = [ + "POLICIES", + "ExecutionUnit", + "NativeStackMergePolicy", + "OptimizationPolicy", + "PolicyRegistry", + "StackMergeUnit", + "StepUnit", + "UnitContext", + "optimize_plan", + "register_policy", +] diff --git a/src/stack_pr/autoland/optimize/base.py b/src/stack_pr/autoland/optimize/base.py new file mode 100644 index 0000000..7b74add --- /dev/null +++ b/src/stack_pr/autoland/optimize/base.py @@ -0,0 +1,178 @@ +"""Execution units, plan optimization policies, and the registry of policies. + +The plan a user writes is a list of independent steps. Before executing it, +the engine turns the rest of the plan (from ``current_step``) into a list of +*execution units*: by default one ``StepUnit`` per step, which just executes +that step. Each registered ``OptimizationPolicy`` that applies to the run's +options may then replace a span of consecutive units with a composite unit +that does something across several steps (e.g. merge a run of PRs as one +GitHub stack). + +Policies are where cross-step knowledge lives: a policy may look at step types +and at the run, while steps stay independent of each other and of policies. + +The optimized plan is execution-time only. It is never persisted: the state +file holds the user's plan, ``current_step`` stays an index into it, and every +execution (fresh, resumed, or replanned) optimizes the plan again from +``current_step``. A unit runs its steps through ``UnitContext.enter`` and +``UnitContext.execute_step``, so each step is still checkpointed and executed +exactly as a ``StepUnit`` would. + +To add a policy, define a class with a ``name``, ``applies(opts)``, and +``optimize(units, run, opts)``, decorate it with ``@register_policy``, and +import its module from ``optimize/__init__.py``. Policies run in registration +order. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, TypeVar + +if TYPE_CHECKING: + from stack_pr import cli + from stack_pr.autoland.model import LandingContext + from stack_pr.autoland.options import AutolandOptions + from stack_pr.autoland.services import Services + from stack_pr.autoland.steps import StepResult + + +@dataclass(frozen=True) +class UnitContext: + """What an execution unit is given to run its steps with.""" + + run: LandingContext + opts: AutolandOptions + common: cli.CommonArgs + services: Services + # Make plan step *index* the current step and checkpoint, so a resume + # starts from it. Call before doing any of the step's work. + enter: Callable[[int], None] + # Execute plan step *index* (after ``enter``) and record the run-level + # facts it reports (its landed SHA). Returns the step's result. + execute_step: Callable[[int], StepResult] + # Save the run's state now. + checkpoint: Callable[[], None] + + +class ExecutionUnit(ABC): + """A span of consecutive plan steps executed together.""" + + @property + @abstractmethod + def indices(self) -> range: + """The plan step indices the unit covers: consecutive and non-empty.""" + + @abstractmethod + def execute(self, uc: UnitContext) -> str: + """Run the unit's steps, in order. Returns the reason the plan must + stop, or "" once every step is done.""" + + def preview(self) -> list[str]: + """Markup lines announcing, before execution starts, what the unit + does beyond executing its steps one by one.""" + return [] + + +@dataclass(frozen=True) +class StepUnit(ExecutionUnit): + """The default unit: plan step ``index``, executed on its own.""" + + index: int + + @property + def indices(self) -> range: + return range(self.index, self.index + 1) + + def execute(self, uc: UnitContext) -> str: + uc.enter(self.index) + return uc.execute_step(self.index).abort_reason + + +class OptimizationPolicy(Protocol): + """Rewrites the execution plan, replacing spans of units with composites.""" + + # Unique among registered policies. + name: str + + def applies(self, opts: AutolandOptions) -> bool: + """Whether the policy is on for a run with *opts*.""" + ... + + def optimize( + self, + units: list[ExecutionUnit], + run: LandingContext, + opts: AutolandOptions, + ) -> list[ExecutionUnit]: + """The units to execute instead of *units*. + + The result must cover the same plan steps, in order, each exactly once. + Must not change *run*. + """ + ... + + +_P = TypeVar("_P", bound=OptimizationPolicy) + + +class PolicyRegistry: + """The optimization policies, in the order they run.""" + + def __init__(self) -> None: + self._policies: dict[str, OptimizationPolicy] = {} + + def register(self, cls: type[_P]) -> type[_P]: + """Add an instance of *cls* after the policies already registered. + + Raises ValueError if a policy with its name is already registered. + """ + policy = cls() + if policy.name in self._policies: + raise ValueError( + f"Optimization policy {policy.name!r} is already registered by " + f"{type(self._policies[policy.name]).__name__}" + ) + self._policies[policy.name] = policy + return cls + + def policies(self) -> list[OptimizationPolicy]: + """Every registered policy, in registration order.""" + return list(self._policies.values()) + + +POLICIES = PolicyRegistry() + + +def register_policy(cls: type[_P]) -> type[_P]: + """Class decorator adding an optimization policy to the registry.""" + return POLICIES.register(cls) + + +def optimize_plan( + run: LandingContext, + opts: AutolandOptions, + policies: PolicyRegistry | None = None, +) -> list[ExecutionUnit]: + """The execution units for the rest of the plan, from ``run.current_step``. + + Starts from one ``StepUnit`` per step and applies every policy that applies + to *opts*, in registration order. Raises ValueError if a policy returns + units that don't cover the plan's remaining steps in order. + """ + if policies is None: + policies = POLICIES + steps = range(run.current_step, len(run.plan)) + units: list[ExecutionUnit] = [StepUnit(i) for i in steps] + for policy in policies.policies(): + if policy.applies(opts): + units = policy.optimize(list(units), run, opts) + covered = [i for unit in units for i in unit.indices] + if covered != list(steps) or not all(unit.indices for unit in units): + raise ValueError( + f"Optimization policy {policy.name!r} returned units that " + "do not cover the remaining plan steps in order" + ) + return units diff --git a/src/stack_pr/autoland/optimize/stack_merge.py b/src/stack_pr/autoland/optimize/stack_merge.py new file mode 100644 index 0000000..97c13cf --- /dev/null +++ b/src/stack_pr/autoland/optimize/stack_merge.py @@ -0,0 +1,157 @@ +"""``--merge-as-stack``: land runs of consecutive land steps with one stack merge. + +GitHub's mechanics of a stack merge live in ``autoland.native_stack``; this +module decides which steps it covers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from stack_pr.autoland.native_stack import land_as_native_stack +from stack_pr.autoland.optimize.base import ( + ExecutionUnit, + StepUnit, + UnitContext, + register_policy, +) +from stack_pr.autoland.pr_state import PRState +from stack_pr.autoland.steps import LandStep + +if TYPE_CHECKING: + from collections.abc import Iterable + + from stack_pr.autoland.model import LandingContext, StackEntry + from stack_pr.autoland.options import AutolandOptions + + +def _open_land_entry(run: LandingContext, index: int) -> StackEntry | None: + """The still-open PR plan step *index* lands, if it is a land step.""" + step = run.plan[index] + if not isinstance(step, LandStep) or step.already_landed: + return None + entry = run.stack[step.entry_index] + return None if entry.state == PRState.MERGED else entry + + +def stack_merge_run(run: LandingContext, indices: Iterable[int]) -> list[StackEntry]: + """The PRs a stack merge starting at the first of *indices* would land. + + That is the run of consecutive land steps, from the first of *indices*, + whose PRs are still open. Any other step ends the run: a ``w`` or ``c`` + step has to see the PRs above it land separately from the ones below it. + A run of fewer than two PRs gains nothing from a stack merge, so it comes + back empty. + """ + entries: list[StackEntry] = [] + for index in indices: + entry = _open_land_entry(run, index) + if entry is None: + break + entries.append(entry) + return entries if len(entries) > 1 else [] + + +@dataclass(frozen=True) +class StackMergeUnit(ExecutionUnit): + """A run of consecutive open land steps, landed with one stack merge. + + ``entries`` are the run's PRs when the plan was optimized, for the + preview. The unit re-validates the run when it executes: before each of + its steps (until a merge falls through), it looks again at which of the + following steps in the unit still land an open PR, and stack-merges those, + if there are at least two. So a PR merged elsewhere in the meantime + shortens the run, exactly as if the run were evaluated when its first step + is reached. + + After a stack merge, each step still executes on its own, so it reports + its PR as merged, with its landed SHA. If the merge fell through, the steps + it covered land one at a time without retrying the merge. + """ + + span: range + entries: tuple[StackEntry, ...] + + @property + def indices(self) -> range: + return self.span + + def execute(self, uc: UnitContext) -> str: + # The last step of a run whose stack merge fell through. + unstacked_through = -1 + for index in self.span: + uc.enter(index) + if index > unstacked_through: + entries = stack_merge_run(uc.run, range(index, self.span.stop)) + if entries: + result = land_as_native_stack( + entries, ctx=uc.run, common=uc.common, opts=uc.opts + ) + if result.abort_reason: + return result.abort_reason + if not result.landed: + unstacked_through = index + len(entries) - 1 + uc.checkpoint() + abort_reason = uc.execute_step(index).abort_reason + if abort_reason: + return abort_reason + return "" + + def preview(self) -> list[str]: + prs = ", ".join(f"#{e.pr_number}" for e in self.entries) + return [ + ( + f"[dim]Steps {self.span.start + 1}-{self.span.stop} ({prs}) will " + "merge together as a GitHub stack.[/dim]" + ) + ] + + +@register_policy +class NativeStackMergePolicy: + """Group each run of two or more consecutive open land steps into a + ``StackMergeUnit``. On with ``--merge-as-stack``.""" + + name = "native-stack-merge" + + def applies(self, opts: AutolandOptions) -> bool: + return opts.merge_as_stack + + def optimize( + self, + units: list[ExecutionUnit], + run: LandingContext, + opts: AutolandOptions, # noqa: ARG002 - part of the policy interface + ) -> list[ExecutionUnit]: + optimized: list[ExecutionUnit] = [] + # The StepUnits of the run of open land steps being collected. + pending: list[StepUnit] = [] + + def flush() -> None: + if not pending: + return + first = pending[0].index + entries = stack_merge_run(run, range(first, first + len(pending))) + if entries: + optimized.append( + StackMergeUnit( + span=range(first, first + len(entries)), + entries=tuple(entries), + ) + ) + else: + optimized.extend(pending) + pending.clear() + + for unit in units: + if ( + isinstance(unit, StepUnit) + and _open_land_entry(run, unit.index) is not None + ): + pending.append(unit) + else: + flush() + optimized.append(unit) + flush() + return optimized diff --git a/tests/test_autoland.py b/tests/test_autoland.py index 813421b..6620ece 100644 --- a/tests/test_autoland.py +++ b/tests/test_autoland.py @@ -19,7 +19,6 @@ display, engine, gh, - native_stack, options, runtime, waits, @@ -1988,29 +1987,6 @@ def test_merge_as_stack_defaults_on_and_flag_overrides_config() -> None: assert flagged.merge_as_stack is True -def test_native_stack_merge_runs_are_split_by_checkpoints() -> None: - stack = _pinned_stack([101, 102, 103, 104, 105]) - plan = parse_plan("l\nl\nw deploy.yaml\nl\nc QA\nl\nl\n", stack) - ctx = LandingContext(stack=stack, plan=plan) - - runs = native_stack.native_stack_runs(ctx) - - # A lone 'l' between two checkpoints has nothing to merge alongside. - assert [(first, [e.pr_number for e in es]) for first, es in runs] == [ - (0, [101, 102]), - (5, [104, 105]), - ] - - -def test_native_stack_merge_run_starts_above_merged_prs() -> None: - stack = _pinned_stack([101, 102, 103]) - ctx = LandingContext(stack=stack, plan=parse_plan("l\nl\nl\n", stack)) - stack[0].state = PRState.MERGED - - assert native_stack.native_stack_run(ctx, 0) == [] - assert [e.pr_number for e in native_stack.native_stack_run(ctx, 1)] == [102, 103] - - class _FakeNativeStackGitHub: """The slice of GitHub a stack merge talks to, holding the stacks in memory.""" diff --git a/tests/test_autoland_optimize.py b/tests/test_autoland_optimize.py new file mode 100644 index 0000000..f4f9e91 --- /dev/null +++ b/tests/test_autoland_optimize.py @@ -0,0 +1,411 @@ +"""Plan optimization: the policy registry, the native-stack merge policy, and +the units policies build.""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from stack_pr.autoland import engine +from stack_pr.autoland.model import ( + ConfirmStep, + LandingContext, + LandStep, + PRState, + WorkflowStep, +) +from stack_pr.autoland.native_stack import NativeStackMergeResult +from stack_pr.autoland.optimize import ( + ExecutionUnit, + PolicyRegistry, + StackMergeUnit, + StepUnit, + UnitContext, + optimize_plan, +) +from stack_pr.autoland.options import AutolandOptions +from stack_pr.autoland.state import AutolandCheckpointer +from stack_pr.autoland.steps import AutolandStep, StepResult +from tests.helpers import common_args +from tests.test_autoland_step_execution import OPTS, FakeUI, services, stack + +STACKED = dataclasses.replace(OPTS, merge_as_stack=True) + + +def land(*indices: int) -> list[AutolandStep]: + return [LandStep(entry_index=i) for i in indices] + + +def spans(units: list[ExecutionUnit]) -> list[tuple[str, int, int]]: + """Each unit as (type, first step, last step).""" + return [(type(u).__name__, u.indices.start, u.indices.stop - 1) for u in units] + + +# --------------------------------------------------------------------------- +# The policy registry +# --------------------------------------------------------------------------- + + +def _named(policy_name: str) -> type: + class Policy: + name = policy_name + + def applies(self, opts: AutolandOptions) -> bool: + return True + + def optimize( + self, + units: list[ExecutionUnit], + run: LandingContext, + opts: AutolandOptions, + ) -> list[ExecutionUnit]: + return units + + return Policy + + +def test_policies_run_in_registration_order() -> None: + registry = PolicyRegistry() + registry.register(_named("b")) + registry.register(_named("a")) + + assert [p.name for p in registry.policies()] == ["b", "a"] + + +def test_registering_a_policy_name_twice_is_an_error() -> None: + registry = PolicyRegistry() + registry.register(_named("a")) + + with pytest.raises(ValueError, match="'a' is already registered"): + registry.register(_named("a")) + + +def test_a_policy_must_keep_every_step() -> None: + class DropsSteps: + name = "drops-steps" + + def applies(self, opts: AutolandOptions) -> bool: + return True + + def optimize( + self, + units: list[ExecutionUnit], + run: LandingContext, + opts: AutolandOptions, + ) -> list[ExecutionUnit]: + return units[1:] + + registry = PolicyRegistry() + registry.register(DropsSteps) + run = LandingContext(stack=stack(101, 102), plan=land(0, 1)) + + with pytest.raises(ValueError, match="'drops-steps'"): + optimize_plan(run, OPTS, registry) + + +# --------------------------------------------------------------------------- +# The native-stack merge policy +# --------------------------------------------------------------------------- + + +def test_land_runs_are_split_by_workflow_and_confirm_steps() -> None: + plan = [ + *land(0, 1), + WorkflowStep(workflow="deploy.yaml"), + *land(2), + ConfirmStep(condition="QA"), + *land(3, 4), + ] + run = LandingContext(stack=stack(101, 102, 103, 104, 105), plan=plan) + + units = optimize_plan(run, STACKED) + + # A lone 'l' between two checkpoints has nothing to merge alongside. + assert spans(units) == [ + ("StackMergeUnit", 0, 1), + ("StepUnit", 2, 2), + ("StepUnit", 3, 3), + ("StepUnit", 4, 4), + ("StackMergeUnit", 5, 6), + ] + + +def test_already_landed_and_merged_prs_are_left_out_of_runs() -> None: + entries = stack(101, 102, 103, 104, 105) + entries[0].state = PRState.MERGED + entries[3].state = PRState.MERGED + plan = [LandStep(entry_index=-1, pr_number=100), *land(0, 1, 2, 3, 4)] + run = LandingContext(stack=entries, plan=plan) + + units = optimize_plan(run, STACKED) + + assert spans(units) == [ + ("StepUnit", 0, 0), + ("StepUnit", 1, 1), + ("StackMergeUnit", 2, 3), + ("StepUnit", 4, 4), + ("StepUnit", 5, 5), + ] + + +def test_runs_start_at_the_current_step() -> None: + run = LandingContext(stack=stack(101, 102, 103), plan=land(0, 1, 2)) + run.current_step = 1 + + assert spans(optimize_plan(run, STACKED)) == [("StackMergeUnit", 1, 2)] + + +def test_merge_as_stack_off_leaves_every_step_on_its_own() -> None: + run = LandingContext(stack=stack(101, 102), plan=land(0, 1)) + + assert spans(optimize_plan(run, OPTS)) == [ + ("StepUnit", 0, 0), + ("StepUnit", 1, 1), + ] + + +def test_preview_announces_each_stack_merge(mocker) -> None: # noqa: ANN001 + console = mocker.patch("stack_pr.autoland.runtime.console") + plan = [*land(0, 1), WorkflowStep(workflow="deploy.yaml"), *land(2, 3)] + run = LandingContext(stack=stack(101, 102, 103, 104), plan=plan) + + engine.print_execution_preview(run, STACKED) + + assert [c.args[0] for c in console.print.call_args_list] == [ + "[dim]Steps 1-2 (#101, #102) will merge together as a GitHub stack.[/dim]", + "[dim]Steps 4-5 (#103, #104) will merge together as a GitHub stack.[/dim]", + ] + + +# --------------------------------------------------------------------------- +# Executing a StackMergeUnit +# --------------------------------------------------------------------------- + + +class _Recorder: + """A UnitContext's callbacks, recording what the unit does in order.""" + + def __init__(self, run: LandingContext) -> None: + self.run = run + self.events: list[object] = [] + self.abort_at: int | None = None + + def context(self) -> UnitContext: + return UnitContext( + run=self.run, + opts=STACKED, + common=common_args(), + services=services(), + enter=lambda i: self.events.append(("enter", i)), + execute_step=self.execute_step, + checkpoint=lambda: self.events.append("checkpoint"), + ) + + def execute_step(self, index: int) -> StepResult: + self.events.append(("step", index)) + if index == self.abort_at: + return StepResult.abort(f"step {index} failed") + return StepResult.done() + + +def _stack_merge(mocker, results: list[NativeStackMergeResult]): # noqa: ANN001, ANN202 + """Patch the stack merge to return *results* in turn; return the mock.""" + return mocker.patch( + "stack_pr.autoland.optimize.stack_merge.land_as_native_stack", + side_effect=results, + ) + + +def _merged_prs(merge) -> list[list[int]]: # noqa: ANN001 + return [[e.pr_number for e in c.args[0]] for c in merge.call_args_list] + + +def _unit(run: LandingContext) -> StackMergeUnit: + (unit,) = optimize_plan(run, STACKED) + assert isinstance(unit, StackMergeUnit) + return unit + + +def test_stack_merge_then_each_step_executes(mocker) -> None: # noqa: ANN001 + run = LandingContext(stack=stack(101, 102, 103), plan=land(0, 1, 2)) + + def lands(entries, **_kw) -> NativeStackMergeResult: # noqa: ANN001, ANN003 + for e in entries: + e.state = PRState.MERGED + return NativeStackMergeResult(landed=True) + + merge = mocker.patch( + "stack_pr.autoland.optimize.stack_merge.land_as_native_stack", + side_effect=lands, + ) + recorder = _Recorder(run) + + assert _unit(run).execute(recorder.context()) == "" + + assert _merged_prs(merge) == [[101, 102, 103]] + # Each step is entered (checkpointed as current) before any of its work. + assert recorder.events == [ + ("enter", 0), + "checkpoint", + ("step", 0), + ("enter", 1), + ("step", 1), + ("enter", 2), + ("step", 2), + ] + + +def test_a_stack_merge_that_falls_through_lands_the_rest_one_at_a_time( + mocker, # noqa: ANN001 +) -> None: + run = LandingContext(stack=stack(101, 102, 103), plan=land(0, 1, 2)) + merge = _stack_merge(mocker, [NativeStackMergeResult(landed=False)]) + recorder = _Recorder(run) + + assert _unit(run).execute(recorder.context()) == "" + + # The merge is not retried for the steps it covered. + assert _merged_prs(merge) == [[101, 102, 103]] + assert [e for e in recorder.events if e != "checkpoint"] == [ + ("enter", 0), + ("step", 0), + ("enter", 1), + ("step", 1), + ("enter", 2), + ("step", 2), + ] + + +def test_a_stack_merge_abort_stops_before_the_step(mocker) -> None: # noqa: ANN001 + run = LandingContext(stack=stack(101, 102), plan=land(0, 1)) + _stack_merge( + mocker, [NativeStackMergeResult(abort_reason="Stack merge was aborted")] + ) + recorder = _Recorder(run) + + assert _unit(run).execute(recorder.context()) == "Stack merge was aborted" + assert recorder.events == [("enter", 0)] + + +def test_a_step_abort_stops_the_unit(mocker) -> None: # noqa: ANN001 + run = LandingContext(stack=stack(101, 102, 103), plan=land(0, 1, 2)) + _stack_merge(mocker, [NativeStackMergeResult(landed=False)]) + recorder = _Recorder(run) + recorder.abort_at = 1 + + assert _unit(run).execute(recorder.context()) == "step 1 failed" + assert ("enter", 2) not in recorder.events + + +def test_the_run_is_revalidated_when_the_unit_executes(mocker) -> None: # noqa: ANN001 + run = LandingContext(stack=stack(101, 102, 103), plan=land(0, 1, 2)) + unit = _unit(run) + # Merged elsewhere after the plan was optimized. + run.stack[0].state = PRState.MERGED + merge = _stack_merge(mocker, [NativeStackMergeResult(landed=True)]) + recorder = _Recorder(run) + + assert unit.execute(recorder.context()) == "" + + # Step 0 now only skips; the stack merge starts above it. + assert _merged_prs(merge) == [[102, 103]] + assert recorder.events[:4] == [ + ("enter", 0), + ("step", 0), + ("enter", 1), + "checkpoint", + ] + + +def test_a_run_shortened_to_one_pr_is_not_stack_merged(mocker) -> None: # noqa: ANN001 + run = LandingContext(stack=stack(101, 102), plan=land(0, 1)) + unit = _unit(run) + run.stack[1].state = PRState.MERGED + merge = _stack_merge(mocker, []) + recorder = _Recorder(run) + + assert unit.execute(recorder.context()) == "" + + merge.assert_not_called() + assert recorder.events == [("enter", 0), ("step", 0), ("enter", 1), ("step", 1)] + + +# --------------------------------------------------------------------------- +# Adding a policy +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Pair(ExecutionUnit): + """Two steps run back to back, announced together.""" + + first: int + log: list[str] + + @property + def indices(self) -> range: + return range(self.first, self.first + 2) + + def execute(self, uc: UnitContext) -> str: + self.log.append(f"pair {self.first}") + for index in self.indices: + uc.enter(index) + abort_reason = uc.execute_step(index).abort_reason + if abort_reason: + return abort_reason + return "" + + def preview(self) -> list[str]: + return [f"pair {self.first}"] + + +def test_a_policy_defined_outside_the_package_shapes_execution( + tmp_path: Path, + mocker, # noqa: ANN001 +) -> None: + log: list[str] = [] + registry = PolicyRegistry() + + @registry.register + class PairUp: + name = "pair-up" + + def applies(self, opts: AutolandOptions) -> bool: + return True + + def optimize( + self, + units: list[ExecutionUnit], + run: LandingContext, + opts: AutolandOptions, + ) -> list[ExecutionUnit]: + # Steps 0 and 1 become one unit; the rest stay as they are. + assert isinstance(units[0], StepUnit) + return [_Pair(first=0, log=log), *units[2:]] + + plan: list[AutolandStep] = [ConfirmStep(confirmed=True) for _ in range(3)] + run = LandingContext(stack=stack(101), plan=plan) + checkpointer = AutolandCheckpointer( + path=tmp_path / "state.json", branch="feat", base="main" + ) + saves = mocker.spy(checkpointer, "save") + console = mocker.patch("stack_pr.autoland.runtime.console") + + engine.print_execution_preview(run, OPTS, policies=registry) + assert engine.execute_plan( + run, + common_args(), + OPTS, + checkpointer, + services=services(FakeUI()), + policies=registry, + ) + + assert [c.args[0] for c in console.print.call_args_list] == ["pair 0"] + assert log == ["pair 0"] + assert run.current_step == len(plan) + # One checkpoint as each step is entered, and one at the end. + assert saves.call_count == len(plan) + 1