fix: bound runner shutdown waits

cmd/github-actions-runner-job/main.go

@@ -29,6 +29,7 @@ const (
 
 var githubJobPollInterval = 5 * time.Second
 var githubJobExitGracePolls = 3
+var githubRunnerShutdownGrace = 10 * time.Second
 
 func main() {
 	if err := run(); err != nil {
@@ -323,7 +324,7 @@ func (d *runnerDriver) RunGitHubJob(ctx context.Context, mode internal.Ephemeral
 		dispatchedAfter := time.Now().UTC().Add(-10 * time.Second)
 		if err := d.sidecar.dispatchWorkflow(ctx, d.req.Repository, d.req.Workflow, dispatchRef, dispatchInputs); err != nil {
 			runner.cancel()
-			_ = runner.wait()
+			_, _ = runner.waitAfterCancel(githubRunnerShutdownGrace)
 			return result, err
 		}
 		completion, err := d.waitForGitHubCompletion(ctx, runner, spec.RunnerName, dispatchedAfter)
@@ -396,7 +397,7 @@ func (d *runnerDriver) waitForGitHubCompletion(ctx context.Context, runner *runn
 			return last, nil
 		case result := <-runner.result:
 			runner.cancel()
-			err := <-runner.done
+			err, _ := runner.waitAfterCancel(githubRunnerShutdownGrace)
 			if result.success {
 				if completion, observeErr := d.observeGitHubJob(ctx, runnerName, dispatchedAfter); observeErr == nil && completion.WorkflowRunID != 0 {
 					if completion.Terminal && !completion.Success {
@@ -422,18 +423,18 @@ func (d *runnerDriver) waitForGitHubCompletion(ctx context.Context, runner *runn
 				continue
 			}
 			runner.cancel()
-			_ = <-runner.done
+			_, _ = runner.waitAfterCancel(githubRunnerShutdownGrace)
 			if completion.Success {
 				return completion, nil
 			}
 			return completion, fmt.Errorf("github workflow job failed: %s", completion.Message)
 		case <-ctx.Done():
 			runner.cancel()
-			err := <-runner.done
+			err, _ := runner.waitAfterCancel(githubRunnerShutdownGrace)
 			if err != nil {
 				return last, fmt.Errorf("%s timed out waiting for GitHub job completion: %w: %s", filepath.Base(runner.path), ctx.Err(), runner.output())
 			}
-			return last, ctx.Err()
+			return last, fmt.Errorf("%s timed out waiting for GitHub job completion: %w", filepath.Base(runner.path), ctx.Err())
 		}
 	}
 }
@@ -621,6 +622,20 @@ func (r *runningCommand) wait() error {
 	return nil
 }
 
+func (r *runningCommand) waitAfterCancel(timeout time.Duration) (error, bool) {
+	if timeout <= 0 {
+		return <-r.done, true
+	}
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+	select {
+	case err := <-r.done:
+		return err, true
+	case <-timer.C:
+		return nil, false
+	}
+}
+
 func (r *runningCommand) waitForGitHubJob(ctx context.Context) error {
 	select {
 	case err := <-r.done:

cmd/github-actions-runner-job/main_test.go

@@ -2,14 +2,18 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/GoCodeAlone/workflow-plugin-github/internal"
 )
 
 func TestT915CommandRunsDynamicProviderEnvelopeThroughSidecarAndRunner(t *testing.T) {
@@ -332,6 +336,125 @@ func TestT915CommandTreatsGitHubJobAPIAsCompletion(t *testing.T) {
 	}
 }
 
+func TestT915WaitForGitHubCompletionDoesNotBlockOnRunnerShutdownAfterTerminalSuccess(t *testing.T) {
+	oldPoll := githubJobPollInterval
+	githubJobPollInterval = 10 * time.Millisecond
+	oldShutdownGrace := githubRunnerShutdownGrace
+	githubRunnerShutdownGrace = 10 * time.Millisecond
+	t.Cleanup(func() {
+		githubJobPollInterval = oldPoll
+		githubRunnerShutdownGrace = oldShutdownGrace
+	})
+
+	var canceled bool
+	sidecar := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.Header.Get("Authorization"); got != "Bearer provider-token" {
+			t.Fatalf("authorization = %q", got)
+		}
+		switch {
+		case r.Method == http.MethodGet && r.URL.Path == "/v1/actions/repos/GoCodeAlone/workflow-compute/workflows/dogfood.yml/runs":
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(`{"workflow_runs":[{"id":28466630619,"status":"completed","conclusion":"success"}]}`))
+		case r.Method == http.MethodGet && r.URL.Path == "/v1/actions/repos/GoCodeAlone/workflow-compute/actions/runs/28466630619/jobs":
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(`{"jobs":[{"id":84367972241,"run_id":28466630619,"status":"completed","conclusion":"success","runner_name":"wfc-stg-ghp-linux-260629fabb6f-1820455d6b7c"}]}`))
+		default:
+			t.Fatalf("unexpected sidecar request: %s %s", r.Method, r.URL.String())
+		}
+	}))
+	t.Cleanup(sidecar.Close)
+
+	driver := &runnerDriver{
+		req: internal.EphemeralRunnerJobRequest{
+			Repository: "GoCodeAlone/workflow-compute",
+			Workflow:   "dogfood.yml",
+		},
+		sidecar: &providerSidecarClient{
+			baseURL: sidecar.URL,
+			token:   "provider-token",
+			http:    sidecar.Client(),
+		},
+	}
+	runner := &runningCommand{
+		path: "run.sh",
+		cancel: func() {
+			canceled = true
+		},
+		result: make(chan runnerCompletion, 1),
+		done:   make(chan error),
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() {
+		completion, err := driver.waitForGitHubCompletion(ctx, runner, "wfc-stg-ghp-linux-260629fabb6f-1820455d6b7c", time.Now().UTC().Add(-time.Minute))
+		if err != nil {
+			done <- err
+			return
+		}
+		if !completion.Success || !completion.Terminal || completion.WorkflowRunID != 28466630619 || completion.WorkflowJobID != 84367972241 {
+			done <- fmt.Errorf("completion = %+v", completion)
+			return
+		}
+		done <- nil
+	}()
+
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-time.After(150 * time.Millisecond):
+		t.Fatal("waitForGitHubCompletion blocked on runner shutdown after terminal success")
+	}
+	if !canceled {
+		t.Fatal("runner was not asked to stop")
+	}
+}
+
+func TestT915WaitForGitHubCompletionDoesNotBlockOnRunnerShutdownAfterTimeout(t *testing.T) {
+	oldPoll := githubJobPollInterval
+	githubJobPollInterval = time.Hour
+	oldShutdownGrace := githubRunnerShutdownGrace
+	githubRunnerShutdownGrace = 10 * time.Millisecond
+	t.Cleanup(func() {
+		githubJobPollInterval = oldPoll
+		githubRunnerShutdownGrace = oldShutdownGrace
+	})
+
+	var canceled bool
+	driver := &runnerDriver{}
+	runner := &runningCommand{
+		path: "run.sh",
+		cancel: func() {
+			canceled = true
+		},
+		result: make(chan runnerCompletion, 1),
+		done:   make(chan error),
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() {
+		_, err := driver.waitForGitHubCompletion(ctx, runner, "wfc-stg-ghp-linux-timeout", time.Now().UTC().Add(-time.Minute))
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		if err == nil || !strings.Contains(err.Error(), "timed out waiting for GitHub job completion") {
+			t.Fatalf("err = %v, want timeout", err)
+		}
+	case <-time.After(150 * time.Millisecond):
+		t.Fatal("waitForGitHubCompletion blocked on runner shutdown after timeout")
+	}
+	if !canceled {
+		t.Fatal("runner was not asked to stop")
+	}
+}
+
 func TestT915CommandRejectsFailedGitHubJobAPICompletion(t *testing.T) {
 	oldPoll := githubJobPollInterval
 	githubJobPollInterval = 10 * time.Millisecond