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