Auto merge of #124743 - Kobzol:arbitrary-try-build, r=pietroalbini

CI: enable arbitrary try builds, take two

Fixed version of https://github.com/rust-lang/rust/pull/124631, which hopefully won't completely break our CI this time 🤦‍♂️ Sorry once again. Only the last commit is new.

r? `@pietroalbini`

src/ci/github-actions/calculate-job-matrix.py

@@ -8,10 +8,11 @@ It reads job definitions from `src/ci/github-actions/jobs.yml`
 and filters them based on the event that happened on CI.
 """
 import dataclasses
-import enum
 import json
 import logging
 import os
+import re
+import typing
 from pathlib import Path
 from typing import List, Dict, Any, Optional
 
@@ -44,10 +45,22 @@ def add_base_env(jobs: List[Job], environment: Dict[str, str]) -> List[Job]:
     return jobs
 
 
-class WorkflowRunType(enum.Enum):
-    PR = enum.auto()
-    Try = enum.auto()
-    Auto = enum.auto()
+@dataclasses.dataclass
+class PRRunType:
+    pass
+
+
+@dataclasses.dataclass
+class TryRunType:
+    custom_jobs: List[str]
+
+
+@dataclasses.dataclass
+class AutoRunType:
+    pass
+
+
+WorkflowRunType = typing.Union[PRRunType, TryRunType, AutoRunType]
 
 
 @dataclasses.dataclass
@@ -55,11 +68,28 @@ class GitHubCtx:
     event_name: str
     ref: str
     repository: str
+    commit_message: Optional[str]
+
+
+def get_custom_jobs(ctx: GitHubCtx) -> List[str]:
+    """
+    Tries to parse names of specific CI jobs that should be executed in the form of
+    try-job: <job-name>
+    from the commit message of the passed GitHub context.
+    """
+    if ctx.commit_message is None:
+        return []
+
+    regex = re.compile(r"^try-job: (.*)", re.MULTILINE)
+    jobs = []
+    for match in regex.finditer(ctx.commit_message):
+        jobs.append(match.group(1))
+    return jobs
 
 
 def find_run_type(ctx: GitHubCtx) -> Optional[WorkflowRunType]:
     if ctx.event_name == "pull_request":
-        return WorkflowRunType.PR
+        return PRRunType()
     elif ctx.event_name == "push":
         old_bors_try_build = (
             ctx.ref in ("refs/heads/try", "refs/heads/try-perf") and
@@ -72,20 +102,41 @@ def find_run_type(ctx: GitHubCtx) -> Optional[WorkflowRunType]:
         try_build = old_bors_try_build or new_bors_try_build
 
         if try_build:
-            return WorkflowRunType.Try
+            jobs = get_custom_jobs(ctx)
+            return TryRunType(custom_jobs=jobs)
 
         if ctx.ref == "refs/heads/auto" and ctx.repository == "rust-lang-ci/rust":
-            return WorkflowRunType.Auto
+            return AutoRunType()
 
     return None
 
 
 def calculate_jobs(run_type: WorkflowRunType, job_data: Dict[str, Any]) -> List[Job]:
-    if run_type == WorkflowRunType.PR:
+    if isinstance(run_type, PRRunType):
         return add_base_env(name_jobs(job_data["pr"], "PR"), job_data["envs"]["pr"])
-    elif run_type == WorkflowRunType.Try:
-        return add_base_env(name_jobs(job_data["try"], "try"), job_data["envs"]["try"])
-    elif run_type == WorkflowRunType.Auto:
+    elif isinstance(run_type, TryRunType):
+        jobs = job_data["try"]
+        custom_jobs = run_type.custom_jobs
+        if custom_jobs:
+            if len(custom_jobs) > 10:
+                raise Exception(
+                    f"It is only possible to schedule up to 10 custom jobs, "
+                    f"received {len(custom_jobs)} jobs"
+                )
+
+            jobs = []
+            unknown_jobs = []
+            for custom_job in custom_jobs:
+                job = [j for j in job_data["auto"] if j["image"] == custom_job]
+                if not job:
+                    unknown_jobs.append(custom_job)
+                    continue
+                jobs.append(job[0])
+            if unknown_jobs:
+                raise Exception(f"Custom job(s) `{unknown_jobs}` not found in auto jobs")
+
+        return add_base_env(name_jobs(jobs, "try"), job_data["envs"]["try"])
+    elif isinstance(run_type, AutoRunType):
         return add_base_env(name_jobs(job_data["auto"], "auto"), job_data["envs"]["auto"])
 
     return []
@@ -99,19 +150,25 @@ def skip_jobs(jobs: List[Dict[str, Any]], channel: str) -> List[Job]:
 
 
 def get_github_ctx() -> GitHubCtx:
+    event_name = os.environ["GITHUB_EVENT_NAME"]
+
+    commit_message = None
+    if event_name == "push":
+        commit_message = os.environ["COMMIT_MESSAGE"]
     return GitHubCtx(
-        event_name=os.environ["GITHUB_EVENT_NAME"],
+        event_name=event_name,
         ref=os.environ["GITHUB_REF"],
-        repository=os.environ["GITHUB_REPOSITORY"]
+        repository=os.environ["GITHUB_REPOSITORY"],
+        commit_message=commit_message
     )
 
 
 def format_run_type(run_type: WorkflowRunType) -> str:
-    if run_type == WorkflowRunType.PR:
+    if isinstance(run_type, PRRunType):
         return "pr"
-    elif run_type == WorkflowRunType.Auto:
+    elif isinstance(run_type, AutoRunType):
         return "auto"
-    elif run_type == WorkflowRunType.Try:
+    elif isinstance(run_type, TryRunType):
         return "try"
     else:
         raise AssertionError()
@@ -135,6 +192,10 @@ if __name__ == "__main__":
     if run_type is not None:
         jobs = calculate_jobs(run_type, data)
     jobs = skip_jobs(jobs, channel)
+
+    if not jobs:
+        raise Exception("Scheduled job list is empty, this is an error")
+
     run_type = format_run_type(run_type)
 
     logging.info(f"Output:\n{yaml.dump(dict(jobs=jobs, run_type=run_type), indent=4)}")