`4fc3462`

runner: poll cancellation and stop active containers

Authored by mfwolffe <wolffemf@dukes.jmu.edu> 2 days ago

SHA: 4fc3462a651363c8271459699365a4dbffe820cb
Parents: 105eb12
Tree: 1f7caae

5 changed files

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 {
  type SleepFunc func(ctx context.Context, d time.Duration) error
  type Options struct {
 -	API          API
 -	Engine       engine.Engine
 -	Workspaces   Workspaces
 -	Logger       *slog.Logger
 -	Labels       []string
 -	Capacity     int
 -	PollInterval time.Duration
 -	DefaultImage string
 -	Clock        func() time.Time
 -	Sleep        SleepFunc
 +	API                API
 +	Engine             engine.Engine
 +	Workspaces         Workspaces
 +	Logger             *slog.Logger
 +	Labels             []string
 +	Capacity           int
 +	PollInterval       time.Duration
 +	CancelPollInterval time.Duration
 +	DefaultImage       string
 +	Clock              func() time.Time
 +	Sleep              SleepFunc
+ }
  type Runner struct {
 -	api          API
 -	engine       engine.Engine
 -	workspaces   Workspaces
 -	logger       *slog.Logger
 -	labels       []string
 -	capacity     int
 -	pollInterval time.Duration
 -	defaultImage string
 -	clock        func() time.Time
 -	sleep        SleepFunc
 +	api                API
 +	engine             engine.Engine
 +	workspaces         Workspaces
 +	logger             *slog.Logger
 +	labels             []string
 +	capacity           int
 +	pollInterval       time.Duration
 +	cancelPollInterval time.Duration
 +	defaultImage       string
 +	clock              func() time.Time
 +	sleep              SleepFunc
+ }
  func New(opts Options) *Runner {
  	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
+ 	}
  	return &Runner{
 -		api:          opts.API,
 -		engine:       opts.Engine,
 -		workspaces:   opts.Workspaces,
 -		logger:       logger,
 -		labels:       append([]string{}, opts.Labels...),
 -		capacity:     capacity,
 -		pollInterval: poll,
 -		defaultImage: opts.DefaultImage,
 -		clock:        clock,
 -		sleep:        sleep,
 +		api:                opts.API,
 +		engine:             opts.Engine,
 +		workspaces:         opts.Workspaces,
 +		logger:             logger,
 +		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)
+ }
  type fakeEngine struct {
 -	job  engine.Job
 -	out  engine.Outcome
 -	logs []engine.LogChunk
 -	err  error
 +	job       engine.Job
 +	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}}}