Restart cancelled workflows on vm preemption

src/lib/githubClient/githubClient.go

@@ -66,6 +66,7 @@ type WorkflowRunInfo struct {
 	Conclusion string
 	Event      string
 	HeadBranch string
+	WorkflowID int64
 	CreatedAt  time.Time
 }
 
@@ -82,10 +83,38 @@ func (gc *GithubClient) GetWorkflowRunInfo(repoName string, workflowRunId int64)
 		Conclusion: wr.GetConclusion(),
 		Event:      wr.GetEvent(),
 		HeadBranch: wr.GetHeadBranch(),
+		WorkflowID: wr.GetWorkflowID(),
 		CreatedAt:  wr.GetCreatedAt().Time,
 	}, nil
 }
 
+// HasNewerWorkflowRun reports whether the workflow has a run for the same
+// branch and event created after the given run. Run IDs are monotonically
+// increasing, so the newest run having a higher ID means the given run has
+// been superseded.
+func (gc *GithubClient) HasNewerWorkflowRun(repoName string, workflowID int64, branch string, event string, runID int64) (bool, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), githubAPITimeout)
+	defer cancel()
+
+	runs, _, err := gc.client.Actions.ListWorkflowRunsByID(ctx, gc.Org.Name, repoName, workflowID, &github.ListWorkflowRunsOptions{
+		Branch:              branch,
+		Event:               event,
+		ExcludePullRequests: true,
+		ListOptions:         github.ListOptions{PerPage: 1},
+	})
+	if err != nil {
+		return false, err
+	}
+
+	// Runs are listed newest-first.
+	for _, run := range runs.WorkflowRuns {
+		if run.GetID() > runID {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 // HasClosedPullRequestForBranch reports whether a pull request with the given
 // head branch was closed (merged or not) after the given time, while no pull
 // request for that branch is currently open. Detection is by head branch

src/lib/restarter/workflowRestarter.go

@@ -92,7 +92,11 @@ func (wr *WorkflowRestarter) handleRestartRequest(ctx context.Context, logger *l
 	}
 
 	switch run.Conclusion {
-	case "failure":
+	// "cancelled" is included because GitHub marks some jobs on a preempted
+	// runner as cancelled rather than failed, and a single cancelled job makes
+	// the whole run conclude "cancelled" — such runs are preemption victims
+	// just like failed ones.
+	case "failure", "cancelled":
 		// A failed merge queue run already evicted the PR from the queue;
 		// re-running it cannot re-enqueue the PR, so a restart is wasted.
 		if run.Event == "merge_group" {
@@ -113,6 +117,20 @@ func (wr *WorkflowRestarter) handleRestartRequest(ctx context.Context, logger *l
 			break
 		}
 
+		if run.Conclusion == "cancelled" {
+			superseded, err := wr.runSuperseded(logger, ghClient, req, run)
+			if err != nil {
+				// Leave the request pending: it is retried on the next poll
+				// and eventually expires via TTL.
+				return
+			}
+			if superseded {
+				logger.Infof("Skipping restart for cancelled workflow run %d (%s/%s): a newer run exists for branch '%s'",
+					req.WorkflowRunId, req.OrgName, req.RepoName, run.HeadBranch)
+				break
+			}
+		}
+
 		logger.Infof("Restarting failed jobs for workflow run %d (%s/%s)", req.WorkflowRunId, req.OrgName, req.RepoName)
 		err = ghClient.RestartFailedJobs(req.RepoName, req.WorkflowRunId)
 		if err != nil {
@@ -129,6 +147,28 @@ func (wr *WorkflowRestarter) handleRestartRequest(ctx context.Context, logger *l
 	}
 }
 
+// runSuperseded reports whether a newer run of the same workflow exists for
+// the run's branch and event. A cancelled run with a newer sibling was almost
+// certainly cancelled by a concurrency group when the newer run started, not
+// by preemption, and restarting it would waste runners on an obsolete commit.
+// A manually cancelled run without a newer sibling is indistinguishable from a
+// preemption victim and gets restarted; that trade-off is accepted since
+// restart requests only exist for runs that lost a runner to preemption.
+func (wr *WorkflowRestarter) runSuperseded(logger *log.Entry, ghClient *githubClient.GithubClient, req repositories.RestartRequest, run githubClient.WorkflowRunInfo) (bool, error) {
+	// Without a branch or workflow id the newer-run lookup cannot be scoped;
+	// fail open and restart.
+	if run.HeadBranch == "" || run.WorkflowID == 0 {
+		return false, nil
+	}
+
+	superseded, err := ghClient.HasNewerWorkflowRun(req.RepoName, run.WorkflowID, run.HeadBranch, run.Event, req.WorkflowRunId)
+	if err != nil {
+		logger.Errorf("Failed to check for newer runs of workflow run %d (branch '%s'): %v", req.WorkflowRunId, run.HeadBranch, err)
+		return false, err
+	}
+	return superseded, nil
+}
+
 // branchPRClosed reports whether the run's head branch belongs to a pull
 // request closed (merged or not) after the run was created. The closedAfter
 // guard keeps an old closed PR from a reused branch name from suppressing a

src/lib/restarter/workflowRestarter_test.go

@@ -61,15 +61,18 @@ var _ repositories.RestarterRepository = (*mockRestarterRepository)(nil)
 
 // --- Fake GitHub API ---
 
-// fakeGithubAPI serves the three endpoints the restarter hits, backed by
+// fakeGithubAPI serves the four endpoints the restarter hits, backed by
 // canned JSON responses.
 type fakeGithubAPI struct {
-	runJSON   string
-	prsJSON   string
-	prsStatus int
-
-	prCalls  int
-	restarts int
+	runJSON       string
+	prsJSON       string
+	prsStatus     int
+	runListJSON   string
+	runListStatus int
+
+	prCalls      int
+	runListCalls int
+	restarts     int
 }
 
 func (f *fakeGithubAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -78,6 +81,13 @@ func (f *fakeGithubAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	case strings.HasSuffix(r.URL.Path, "/rerun-failed-jobs"):
 		f.restarts++
 		w.WriteHeader(http.StatusCreated)
+	case strings.Contains(r.URL.Path, "/actions/workflows/"):
+		f.runListCalls++
+		if f.runListStatus != 0 {
+			w.WriteHeader(f.runListStatus)
+			return
+		}
+		_, _ = w.Write([]byte(f.runListJSON))
 	case strings.Contains(r.URL.Path, "/actions/runs/"):
 		_, _ = w.Write([]byte(f.runJSON))
 	case strings.HasSuffix(r.URL.Path, "/pulls"):
@@ -93,10 +103,18 @@ func (f *fakeGithubAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 func runJSON(status, conclusion, headBranch, event string) string {
-	return fmt.Sprintf(`{"id":42,"status":%q,"conclusion":%q,"head_branch":%q,"event":%q,"created_at":%q}`,
+	return fmt.Sprintf(`{"id":42,"workflow_id":7,"status":%q,"conclusion":%q,"head_branch":%q,"event":%q,"created_at":%q}`,
 		status, conclusion, headBranch, event, time.Now().Add(-10*time.Minute).UTC().Format(time.RFC3339))
 }
 
+func runListJSON(runIDs ...int64) string {
+	var runs []string
+	for _, id := range runIDs {
+		runs = append(runs, fmt.Sprintf(`{"id":%d}`, id))
+	}
+	return fmt.Sprintf(`{"total_count":%d,"workflow_runs":[%s]}`, len(runIDs), strings.Join(runs, ","))
+}
+
 func prListJSON(prs ...string) string {
 	return "[" + strings.Join(prs, ",") + "]"
 }
@@ -329,6 +347,93 @@ func TestHandleRestartRequest_NoHeadBranchStillRestarts(t *testing.T) {
 	assert.Contains(t, repo.deleted, int64(42))
 }
 
+func TestHandleRestartRequest_CancelledRestarts(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON:     runJSON("completed", "cancelled", "main", "push"),
+		runListJSON: runListJSON(42), // the run itself is the newest — no supersession
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Equal(t, 1, api.runListCalls, "cancelled runs must be checked for supersession")
+	assert.Equal(t, 1, api.restarts, "a cancelled run without a newer sibling is a preemption victim and must restart")
+	assert.Contains(t, repo.deleted, int64(42))
+}
+
+func TestHandleRestartRequest_CancelledSupersededSkipsRestart(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON:     runJSON("completed", "cancelled", "feature", "pull_request"),
+		prsJSON:     prListJSON(openPR()),
+		runListJSON: runListJSON(43), // a newer run cancelled this one via concurrency
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Zero(t, api.restarts, "a run superseded by a newer one must not be restarted")
+	assert.Contains(t, repo.deleted, int64(42))
+}
+
+func TestHandleRestartRequest_CancelledSupersededCheckErrorKeepsRequest(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON:       runJSON("completed", "cancelled", "main", "push"),
+		runListStatus: http.StatusInternalServerError,
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Zero(t, api.restarts)
+	assert.Empty(t, repo.deleted, "request must stay pending for retry on supersession-check error")
+}
+
+func TestHandleRestartRequest_CancelledNoHeadBranchSkipsSupersededCheck(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON:     runJSON("completed", "cancelled", "", "push"),
+		runListJSON: runListJSON(43), // must not be consulted without a head branch
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Zero(t, api.runListCalls, "supersession check cannot be scoped without a branch")
+	assert.Equal(t, 1, api.restarts, "fail open and restart when supersession cannot be checked")
+	assert.Contains(t, repo.deleted, int64(42))
+}
+
+func TestHandleRestartRequest_CancelledMergeGroupSkipsRestart(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON: runJSON("completed", "cancelled", "gh-readonly-queue/main/pr-7", "merge_group"),
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Zero(t, api.restarts, "cancelled merge queue runs must not be restarted")
+	assert.Contains(t, repo.deleted, int64(42))
+}
+
+func TestHandleRestartRequest_FailureSkipsSupersededCheck(t *testing.T) {
+	repo := &mockRestarterRepository{}
+	api := &fakeGithubAPI{
+		runJSON:     runJSON("completed", "failure", "main", "push"),
+		runListJSON: runListJSON(43), // must not be consulted for failed runs
+	}
+	wr := newTestRestarter(t, repo, api)
+
+	wr.handleRestartRequest(context.Background(), log.WithField("test", true), testRequest())
+
+	assert.Zero(t, api.runListCalls, "failed runs restart without a supersession check")
+	assert.Equal(t, 1, api.restarts)
+	assert.Contains(t, repo.deleted, int64(42))
+}
+
 func TestHandleRestartRequest_NotCompleted(t *testing.T) {
 	repo := &mockRestarterRepository{}
 	api := &fakeGithubAPI{