runner: poll cancellation and stop active containers

Status	File	+	-
M	`internal/runner/api/client_test.go`	26	0
M	`internal/runner/engine/docker.go`	55	3
M	`internal/runner/engine/docker_test.go`	80	4
M	`internal/runner/runner.go`	114	35
M	`internal/runner/runner_test.go`	97	5

internal/runner/api/client_test.gomodified

  		t.Fatalf("AppendLog: %v", err)
+ 	}
+ }
++
 +func TestCancelCheck_UsesJobTokenAndParsesResponse(t *testing.T) {
 +	t.Parallel()
 +	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 +		if r.URL.Path != "/api/v1/jobs/10/cancel-check" {
 +			t.Fatalf("path: %s", r.URL.Path)
 +		}
 +		if got := r.Header.Get("Authorization"); got != "Bearer job-token" {
 +			t.Fatalf("Authorization: %q", got)
 +		}
 +		w.Header().Set("Content-Type", "application/json")
 +		_, _ = w.Write([]byte(`{"cancelled":true,"next_token":"next","next_token_expires_at":"2026-05-10T21:15:00Z"}`))
 +	}))
 +	defer srv.Close()
 +	client, err := New(Config{BaseURL: srv.URL, RunnerToken: "runner-token", HTTPClient: srv.Client()})
 +	if err != nil {
 +		t.Fatalf("New: %v", err)
 +	}
 +	resp, err := client.CancelCheck(t.Context(), 10, "job-token")
 +	if err != nil {
 +		t.Fatalf("CancelCheck: %v", err)
 +	}
 +	if !resp.Cancelled || resp.NextToken != "next" {
 +		t.Fatalf("response: %#v", resp)
 +	}
 +}

internal/runner/engine/docker.gomodified

  	cfg       DockerConfig
  	streams   map[int64]chan LogChunk
  	eventSubs map[int64]chan Event
 +	active    map[int64]string
  	mu        sync.Mutex
+ }
  	if cfg.Logger == nil {
  		cfg.Logger = slog.New(slog.NewTextHandler(io.Discard, nil))
+ 	}
 -	return &Docker{cfg: cfg, streams: make(map[int64]chan LogChunk), eventSubs: make(map[int64]chan Event)}
 +	return &Docker{
 +		cfg:       cfg,
 +		streams:   make(map[int64]chan LogChunk),
 +		eventSubs: make(map[int64]chan Event),
 +		active:    make(map[int64]string),
 +	}
+ }
  func (d *Docker) Execute(ctx context.Context, job Job) (Outcome, error) {
  	if err != nil {
  		return err
+ 	}
 +	d.setActiveContainer(job.ID, invocation.containerName)
 +	defer d.clearActiveContainer(job.ID, invocation.containerName)
  	d.logStep(ctx, "runner step starting", job, step, invocation, "")
  	writer := d.newStepLogWriter(ctx, job.ID, step.ID, job.MaskValues)
  	out := io.MultiWriter(d.cfg.Stdout, writer)
  type dockerInvocation struct {
  	args           []string
  	env            []string
 +	containerName  string
  	image          string
  	network        string
  	memory         string
  	if d.cfg.AllowRoot && permissionsRequestRoot(job.Permissions) {
  		user = "0:0"
+ 	}
 +	containerName := dockerContainerName(job, step)
  	args := []string{
  		"run",
  		"--rm",
 +		"--name", containerName,
  		"--network=" + d.cfg.Network,
  		"--memory=" + d.cfg.Memory,
  		"--cpus=" + d.cfg.CPUs,
  	return dockerInvocation{
  		args:           args,
  		env:            processEnv,
 +		containerName:  containerName,
  		image:          image,
  		network:        d.cfg.Network,
  		memory:         d.cfg.Memory,
  	return d.ensureEventStream(jobID), nil
+ }
 -func (d *Docker) Cancel(_ context.Context, _ int64) error {
 -	return ErrUnsupported
 +func (d *Docker) Cancel(ctx context.Context, jobID int64) error {
 +	name := d.activeContainer(jobID)
 +	if name == "" {
 +		return nil
 +	}
 +	killCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
 +	defer cancel()
 +	if err := d.cfg.Runner.Run(killCtx, d.cfg.Binary, []string{"kill", name}, nil, d.cfg.Stdout, d.cfg.Stderr); err != nil {
 +		return fmt.Errorf("runner engine: kill container %s: %w", name, err)
 +	}
 +	return nil
 +}
++
 +func (d *Docker) setActiveContainer(jobID int64, name string) {
 +	if name == "" {
 +		return
 +	}
 +	d.mu.Lock()
 +	d.active[jobID] = name
 +	d.mu.Unlock()
 +}
++
 +func (d *Docker) clearActiveContainer(jobID int64, name string) {
 +	d.mu.Lock()
 +	if d.active[jobID] == name {
 +		delete(d.active, jobID)
 +	}
 +	d.mu.Unlock()
 +}
++
 +func (d *Docker) activeContainer(jobID int64) string {
 +	d.mu.Lock()
 +	defer d.mu.Unlock()
 +	return d.active[jobID]
+ }
  func (d *Docker) ensureStream(jobID int64) chan LogChunk {
  	return clean, nil
+ }
 +func dockerContainerName(job Job, step Step) string {
 +	stepID := step.ID
 +	if stepID == 0 {
 +		stepID = int64(step.Index)
 +	}
 +	return fmt.Sprintf("shithub-job-%d-step-%d", job.ID, stepID)
 +}
++
  var envNameRE = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)
  func validateEnv(env map[string]string) (map[string]string, error) {

internal/runner/engine/docker_test.gomodified

  	"io"
  	"reflect"
  	"strings"
 +	"sync"
  	"testing"
  	"time"
+ )
  	return nil
+ }
 +type cancellableRunner struct {
 +	started  chan struct{}
 +	killed   chan struct{}
 +	killArgs []string
 +	mu       sync.Mutex
 +}
++
 +func newCancellableRunner() *cancellableRunner {
 +	return &cancellableRunner{
 +		started: make(chan struct{}),
 +		killed:  make(chan struct{}),
 +	}
 +}
++
 +func (r *cancellableRunner) Run(ctx context.Context, _ string, args []string, _ []string, _, _ io.Writer) error {
 +	if len(args) > 0 && args[0] == "kill" {
 +		r.mu.Lock()
 +		r.killArgs = append([]string{}, args...)
 +		r.mu.Unlock()
 +		close(r.killed)
 +		return nil
 +	}
 +	close(r.started)
 +	select {
 +	case <-r.killed:
 +		return context.Canceled
 +	case <-ctx.Done():
 +		return ctx.Err()
 +	}
 +}
++
  func TestDockerExecute_BuildsResourceCappedRunCommand(t *testing.T) {
  	t.Parallel()
  	rec := &recordingRunner{}
  		t.Fatalf("Conclusion: %q", out.Conclusion)
+ 	}
  	want := []string{
 -		"run", "--rm", "--network=none", "--memory=2g", "--cpus=2",
 +		"run", "--rm", "--name", "shithub-job-1-step-0",
 +		"--network=none", "--memory=2g", "--cpus=2",
  		"--pids-limit=512", "--read-only",
  		"--tmpfs", "/tmp:rw,exec,nosuid,nodev,size=1g",
  		"--cap-drop=ALL", "--cap-add=DAC_OVERRIDE", "--cap-add=SETGID", "--cap-add=SETUID",
  		"--ulimit", "nofile=4096:4096", "--ulimit", "nproc=512:512",
  		"--user", "65534:65534",
  		"--workdir=/workspace/subdir",
 -		"--mount", rec.args[23],
 +		"--mount", rec.args[25],
  		"--env", "A", "--env", "B",
  		"runner-image", "bash", "-c", "echo hi",
+ 	}
  	if !reflect.DeepEqual(rec.args, want) {
  		t.Fatalf("args:\ngot  %#v\nwant %#v", rec.args, want)
+ 	}
 -	if !strings.HasPrefix(rec.args[23], "type=bind,src=") || !strings.HasSuffix(rec.args[23], ",dst=/workspace,rw") {
 -		t.Fatalf("workspace mount arg: %q", rec.args[23])
 +	if !strings.HasPrefix(rec.args[25], "type=bind,src=") || !strings.HasSuffix(rec.args[25], ",dst=/workspace,rw") {
 +		t.Fatalf("workspace mount arg: %q", rec.args[25])
+ 	}
  	if wantEnv := []string{"A=job", "B=step"}; !reflect.DeepEqual(rec.env, wantEnv) {
  		t.Fatalf("env:\ngot  %#v\nwant %#v", rec.env, wantEnv)
+ 	}
+ }
 +func TestDockerCancelKillsActiveContainer(t *testing.T) {
 +	t.Parallel()
 +	rec := newCancellableRunner()
 +	d := NewDocker(DockerConfig{
 +		DefaultImage: "runner-image",
 +		Network:      "bridge",
 +		Memory:       "2g",
 +		CPUs:         "2",
 +		Runner:       rec,
 +	})
 +	type executeResult struct {
 +		out Outcome
 +		err error
 +	}
 +	done := make(chan executeResult, 1)
 +	go func() {
 +		out, err := d.Execute(t.Context(), Job{
 +			ID:           99,
 +			WorkspaceDir: t.TempDir(),
 +			Steps:        []Step{{ID: 123, Run: "sleep 600"}},
 +		})
 +		done <- executeResult{out: out, err: err}
 +	}()
 +	<-rec.started
 +	if err := d.Cancel(t.Context(), 99); err != nil {
 +		t.Fatalf("Cancel: %v", err)
 +	}
 +	res := <-done
 +	if !errors.Is(res.err, context.Canceled) {
 +		t.Fatalf("Execute error: %v", res.err)
 +	}
 +	if res.out.Conclusion != ConclusionCancelled {
 +		t.Fatalf("Conclusion: %q", res.out.Conclusion)
 +	}
 +	rec.mu.Lock()
 +	killArgs := append([]string{}, rec.killArgs...)
 +	rec.mu.Unlock()
 +	want := []string{"kill", "shithub-job-99-step-123"}
 +	if !reflect.DeepEqual(killArgs, want) {
 +		t.Fatalf("kill args: got %#v want %#v", killArgs, want)
 +	}
 +}
++
  func TestDockerExecute_FailureMapsToFailureConclusion(t *testing.T) {
  	t.Parallel()
  	d := NewDocker(DockerConfig{

internal/runner/runner.gomodified

  	"io"
  	"log/slog"
  	"sync"
 +	"sync/atomic"
  	"time"
  	"github.com/tenseleyFlow/shithub/internal/runner/api"
  	UpdateStatus(ctx context.Context, jobID int64, token string, req api.StatusRequest) (api.StatusResponse, error)
  	UpdateStepStatus(ctx context.Context, jobID, stepID int64, token string, req api.StatusRequest) (api.StepStatusResponse, error)
  	AppendLog(ctx context.Context, jobID int64, token string, req api.LogRequest) (api.LogResponse, error)
 +	CancelCheck(ctx context.Context, jobID int64, token string) (api.CancelCheckResponse, error)
+ }
  type Workspaces interface {
  	Labels             []string
  	Capacity           int
  	PollInterval       time.Duration
 +	CancelPollInterval time.Duration
  	DefaultImage       string
  	Clock              func() time.Time
  	Sleep              SleepFunc
  	labels             []string
  	capacity           int
  	pollInterval       time.Duration
 +	cancelPollInterval time.Duration
  	defaultImage       string
  	clock              func() time.Time
  	sleep              SleepFunc
  	if poll <= 0 {
  		poll = 5 * time.Second
+ 	}
 +	cancelPoll := opts.CancelPollInterval
 +	if cancelPoll <= 0 {
 +		cancelPoll = 2 * time.Second
 +	}
  	capacity := opts.Capacity
  	if capacity <= 0 {
  		capacity = 1
  		labels:             append([]string{}, opts.Labels...),
  		capacity:           capacity,
  		pollInterval:       poll,
 +		cancelPollInterval: cancelPoll,
  		defaultImage:       opts.DefaultImage,
  		clock:              clock,
  		sleep:              sleep,
  		}()
+ 	}
 -	outcome, execErr := r.engine.Execute(ctx, toEngineJob(claim.Job, workspaceDir, r.defaultImage))
 +	execCtx, execCancel := context.WithCancel(ctx)
 +	watchCtx, stopCancelWatch := context.WithCancel(ctx)
 +	cancelRequested := atomic.Bool{}
 +	cancelWatchErr := make(chan error, 1)
 +	go func() {
 +		cancelWatchErr <- r.watchCancel(watchCtx, session, claim.Job.ID, execCancel, &cancelRequested)
 +	}()
++
 +	outcome, execErr := r.engine.Execute(execCtx, toEngineJob(claim.Job, workspaceDir, r.defaultImage))
 +	stopCancelWatch()
 +	if err := <-cancelWatchErr; err != nil {
 +		return true, fmt.Errorf("watch job cancellation: %w", err)
 +	}
  	if err := <-drainErr; err != nil {
  		return true, fmt.Errorf("stream runner events: %w", err)
+ 	}
  	if conclusion == "" {
  		conclusion = engine.ConclusionFailure
+ 	}
 +	finalStatus := "completed"
 +	if cancelRequested.Load() {
 +		finalStatus = "cancelled"
 +		conclusion = engine.ConclusionCancelled
 +	}
  	completed := outcome.CompletedAt
  	if completed.IsZero() {
  		completed = r.clock()
  	if outcome.StartedAt.IsZero() {
  		outcome.StartedAt = started
+ 	}
 -	if err := r.complete(ctx, session, conclusion, outcome.StartedAt, completed); err != nil {
 +	if err := r.finish(ctx, session, finalStatus, conclusion, outcome.StartedAt, completed); err != nil {
  		return true, err
+ 	}
 -	if execErr != nil {
 +	if execErr != nil && !cancelRequested.Load() {
  		r.logger.WarnContext(ctx, "job completed with failing engine outcome", "job_id", claim.Job.ID, "conclusion", conclusion, "error", execErr)
+ 	}
  	return true, nil
+ }
  func (r *Runner) complete(ctx context.Context, session *jobSession, conclusion string, started, completed time.Time) error {
 +	return r.finish(ctx, session, "completed", conclusion, started, completed)
 +}
++
 +func (r *Runner) finish(ctx context.Context, session *jobSession, status, conclusion string, started, completed time.Time) error {
  	_, err := session.UpdateStatus(ctx, api.StatusRequest{
 -		Status:      "completed",
 +		Status:      status,
  		Conclusion:  conclusion,
  		StartedAt:   started,
  		CompletedAt: completed,
  	})
  	if err != nil {
 -		return fmt.Errorf("mark job completed: %w", err)
 +		return fmt.Errorf("mark job %s: %w", status, err)
+ 	}
  	return nil
+ }
 +func (r *Runner) watchCancel(
 +	ctx context.Context,
 +	session *jobSession,
 +	jobID int64,
 +	execCancel context.CancelFunc,
 +	cancelRequested *atomic.Bool,
 +) error {
 +	for {
 +		if err := r.sleep(ctx, r.cancelPollInterval); err != nil {
 +			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 +				return nil
 +			}
 +			return err
 +		}
 +		resp, err := session.CancelCheck(ctx)
 +		if err != nil {
 +			if ctx.Err() != nil {
 +				return nil
 +			}
 +			r.logger.WarnContext(ctx, "runner cancel check failed", "job_id", jobID, "error", err)
 +			continue
 +		}
 +		if !resp.Cancelled {
 +			continue
 +		}
 +		cancelRequested.Store(true)
 +		killCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
 +		if err := r.engine.Cancel(killCtx, jobID); err != nil && !errors.Is(err, engine.ErrUnsupported) {
 +			r.logger.WarnContext(ctx, "runner engine cancel failed", "job_id", jobID, "error", err)
 +		}
 +		cancel()
 +		execCancel()
 +		return nil
 +	}
 +}
++
  type jobSession struct {
  	api   API
  	jobID int64
  	return nil
+ }
 +func (s *jobSession) CancelCheck(ctx context.Context) (api.CancelCheckResponse, error) {
 +	s.mu.Lock()
 +	defer s.mu.Unlock()
 +	resp, err := s.api.CancelCheck(ctx, s.jobID, s.token)
 +	if err != nil {
 +		return resp, err
 +	}
 +	if resp.NextToken != "" {
 +		s.token = resp.NextToken
 +	}
 +	return resp, nil
 +}
++
  func drainLogs(ctx context.Context, session *jobSession, logs <-chan engine.LogChunk) error {
  	for {
  		select {

internal/runner/runner_test.gomodified

  	statuses     []api.StatusRequest
  	stepStatuses []api.StatusRequest
  	logs         []api.LogRequest
 +	cancelChecks int
 +	cancelled    bool
  	tokens       []string
  	next         int
+ }
  	return api.LogResponse{NextToken: f.nextToken()}, nil
+ }
 +func (f *fakeAPI) CancelCheck(_ context.Context, _ int64, token string) (api.CancelCheckResponse, error) {
 +	f.tokens = append(f.tokens, token)
 +	f.cancelChecks++
 +	return api.CancelCheckResponse{Cancelled: f.cancelled, NextToken: f.nextToken()}, nil
 +}
++
  func (f *fakeAPI) nextToken() string {
  	f.next++
  	return "next-token-" + strconv.Itoa(f.next)
  	out       engine.Outcome
  	logs      []engine.LogChunk
  	err       error
 +	cancelled bool
+ }
  func (f *fakeEngine) Execute(_ context.Context, job engine.Job) (engine.Outcome, error) {
  	return ch, nil
+ }
 -func (f *fakeEngine) Cancel(_ context.Context, _ int64) error { return nil }
 +func (f *fakeEngine) Cancel(_ context.Context, _ int64) error {
 +	f.cancelled = true
 +	return nil
 +}
  type fakeEventEngine struct {
  	fakeEngine
  	return out, err
+ }
 +type cancelBlockingEngine struct {
 +	job       engine.Job
 +	logs      chan engine.LogChunk
 +	started   chan struct{}
 +	cancelled bool
 +}
++
 +func newCancelBlockingEngine() *cancelBlockingEngine {
 +	return &cancelBlockingEngine{
 +		logs:    make(chan engine.LogChunk),
 +		started: make(chan struct{}),
 +	}
 +}
++
 +func (f *cancelBlockingEngine) Execute(ctx context.Context, job engine.Job) (engine.Outcome, error) {
 +	f.job = job
 +	close(f.started)
 +	<-ctx.Done()
 +	close(f.logs)
 +	now := time.Date(2026, 5, 10, 21, 0, 1, 0, time.UTC)
 +	return engine.Outcome{
 +		Conclusion:  engine.ConclusionCancelled,
 +		StartedAt:   now.Add(-time.Second),
 +		CompletedAt: now,
 +	}, ctx.Err()
 +}
++
 +func (f *cancelBlockingEngine) StreamLogs(_ context.Context, _ int64) (<-chan engine.LogChunk, error) {
 +	return f.logs, nil
 +}
++
 +func (f *cancelBlockingEngine) Cancel(_ context.Context, _ int64) error {
 +	f.cancelled = true
 +	return nil
 +}
++
  type fakeWorkspaces struct {
  	dir     string
  	removed bool
+ 	}
+ }
 +func TestRunOnce_CancelCheckStopsEngineAndMarksJobCancelled(t *testing.T) {
 +	t.Parallel()
 +	now := time.Date(2026, 5, 10, 21, 0, 0, 0, time.UTC)
 +	fapi := &fakeAPI{
 +		claim:     &api.Claim{Token: "job-token", Job: api.Job{ID: 10, RunID: 20}},
 +		cancelled: true,
 +	}
 +	fengine := newCancelBlockingEngine()
 +	r := New(Options{
 +		API:                fapi,
 +		Engine:             fengine,
 +		Workspaces:         &fakeWorkspaces{dir: "/tmp/workspace"},
 +		Clock:              func() time.Time { return now },
 +		CancelPollInterval: time.Nanosecond,
 +		Sleep: func(ctx context.Context, _ time.Duration) error {
 +			select {
 +			case <-ctx.Done():
 +				return ctx.Err()
 +			default:
 +				return nil
 +			}
 +		},
 +	})
 +	claimed, err := r.RunOnce(t.Context())
 +	if err != nil {
 +		t.Fatalf("RunOnce: %v", err)
 +	}
 +	if !claimed {
 +		t.Fatal("claimed = false")
 +	}
 +	if !fengine.cancelled {
 +		t.Fatal("engine Cancel was not called")
 +	}
 +	if fapi.cancelChecks == 0 {
 +		t.Fatal("cancel-check was not called")
 +	}
 +	if len(fapi.statuses) != 2 ||
 +		fapi.statuses[0].Status != "running" ||
 +		fapi.statuses[1].Status != "cancelled" ||
 +		fapi.statuses[1].Conclusion != engine.ConclusionCancelled {
 +		t.Fatalf("statuses: %#v", fapi.statuses)
 +	}
 +}
++
  func TestRunOnce_PrepareFailureMarksJobFailed(t *testing.T) {
  	t.Parallel()
  	fapi := &fakeAPI{claim: &api.Claim{Token: "job-token", Job: api.Job{ID: 10, RunID: 20}}}

tenseleyflow/shithub / `4fc3462`

5 changed files