Skip to content
quarry Artifact [36c04c5bd6]
Login

Artifact [36c04c5bd6]

Artifact 36c04c5bd669692834d08a1d0a93cf9caa70fbb603a56f405511d4df390456b3:


// Package runner is `quarry runner`: a generic CI worker a repository's
// owner runs on their own machine (f24d78d9).
//
// It speaks the engine's HTTP worker protocol through the forge
// (internal/cibridge), so all it needs is the repository's URL and a
// worker token from Settings → Runners. For each task it checks out the
// check-in under test, reads ci/runner.json AT that check-in, runs the
// argv the manifest names for the task's type, streams its output to
// the run's log, and reports the verdict.
//
// It runs whatever the check-in's manifest says. That is the point on
// the owner's machine and exactly why it must never be pointed at a
// repository whose writers the machine's owner does not trust.
package runner

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Checkout puts a check-in's tree in dir and returns nothing else: the
// runner reads the manifest from the tree it is about to test.
type Checkout interface {
	Checkout(ctx context.Context, dir, checkin string) error
}

// Config is what a runner needs.
type Config struct {
	Bridge   string   // https://forge/r/<repo>/ci/bridge — the four worker routes live under it
	Token    string   // a worker token from Settings → Runners
	Tasks    []string // task types to poll for
	MaxTasks int      // tasks run at once
	WorkDir  string   // one workspace per task type under it
	WorkerID string
	Checkout Checkout
	Client   *http.Client
	// Heartbeat is how often a running task extends its claim; the
	// engine redelivers a task whose claim lapses.
	Heartbeat time.Duration
	Logf      func(format string, args ...any)
}

// Runner polls, runs and reports until its context ends.
type Runner struct {
	cfg   Config
	slots chan struct{}
	locks sync.Map // workspace dir -> *sync.Mutex
}

// New checks cfg and fills its defaults.
func New(cfg Config) (*Runner, error) {
	if cfg.Bridge == "" || cfg.Token == "" || len(cfg.Tasks) == 0 || cfg.Checkout == nil || cfg.WorkDir == "" {
		return nil, errors.New("runner: bridge, token, tasks, workdir and checkout are all required")
	}
	if cfg.MaxTasks < 1 {
		cfg.MaxTasks = 1
	}
	if cfg.Client == nil {
		cfg.Client = &http.Client{}
	}
	if cfg.Heartbeat <= 0 {
		cfg.Heartbeat = 20 * time.Second
	}
	if cfg.WorkerID == "" {
		h, _ := os.Hostname()
		cfg.WorkerID = fmt.Sprintf("quarry-runner-%s-%d", h, os.Getpid())
	}
	if cfg.Logf == nil {
		cfg.Logf = log.Printf
	}
	return &Runner{cfg: cfg, slots: make(chan struct{}, cfg.MaxTasks)}, nil
}

// Run connects and polls every task type until ctx ends. One poller per
// type, because a poll's answer does not say which of several types a
// task was for, and the type is what picks the manifest entry.
func (r *Runner) Run(ctx context.Context) error {
	go r.stayConnected(ctx)
	var wg sync.WaitGroup
	for _, t := range r.cfg.Tasks {
		wg.Add(1)
		go func(taskType string) {
			defer wg.Done()
			r.pollLoop(ctx, taskType)
		}(t)
	}
	r.cfg.Logf("quarry runner ready · %s · tasks %s · max-tasks %d", r.cfg.Bridge, strings.Join(r.cfg.Tasks, ", "), r.cfg.MaxTasks)
	wg.Wait()
	return ctx.Err()
}

// stayConnected holds the registration stream open, reconnecting with
// backoff, so the engine lists this worker and its capacity.
func (r *Runner) stayConnected(ctx context.Context) {
	backoff := time.Second
	for ctx.Err() == nil {
		body, _ := json.Marshal(map[string]any{"worker_id": r.cfg.WorkerID, "task_types": r.cfg.Tasks, "max_tasks": r.cfg.MaxTasks})
		resp, err := r.post(ctx, "/v1/workers/connect", body)
		if err == nil {
			backoff = time.Second
			io.Copy(io.Discard, resp.Body) // the heartbeat stream, until it ends
			resp.Body.Close()
		} else if ctx.Err() == nil {
			r.cfg.Logf("connect: %v; retrying in %s", err, backoff)
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(backoff):
		}
		if backoff < time.Minute {
			backoff *= 2
		}
	}
}

type polled struct {
	TaskID string          `json:"task_id"`
	RunID  string          `json:"run_id"`
	StepID string          `json:"step_id"`
	Input  json.RawMessage `json:"input"`
}

// pollLoop takes a slot, then asks for one task of its type; a slot is
// held before claiming so a claimed task never waits for a slot here.
func (r *Runner) pollLoop(ctx context.Context, taskType string) {
	for ctx.Err() == nil {
		select {
		case r.slots <- struct{}{}:
		case <-ctx.Done():
			return
		}
		body, _ := json.Marshal(map[string]any{"task_types": []string{taskType}, "max_tasks": 1, "timeout_ms": 25000})
		resp, err := r.post(ctx, "/v1/tasks/poll", body)
		if err != nil {
			<-r.slots
			if ctx.Err() == nil {
				r.cfg.Logf("poll %s: %v", taskType, err)
				time.Sleep(5 * time.Second)
			}
			continue
		}
		var tasks []polled
		err = json.NewDecoder(resp.Body).Decode(&tasks)
		resp.Body.Close()
		if err != nil || len(tasks) == 0 {
			<-r.slots
			continue
		}
		t := tasks[0]
		go func() {
			defer func() { <-r.slots }()
			r.runTask(ctx, taskType, t)
		}()
	}
}

// runTask checks out, reads the manifest at the check-in, runs the
// task's argv, and resolves the task with the verdict.
func (r *Runner) runTask(ctx context.Context, taskType string, t polled) {
	var in struct {
		Repo    string `json:"repo"`
		Checkin string `json:"checkin"`
	}
	json.Unmarshal(t.Input, &in)
	start := time.Now()
	logs := newLogSink(ctx, r, t.TaskID, "out")
	errs := newLogSink(ctx, r, t.TaskID, "err")
	hbCtx, stopHB := context.WithCancel(ctx)
	go r.heartbeat(hbCtx, t.TaskID)
	verdict := r.execute(ctx, taskType, in.Checkin, logs, errs)
	stopHB()
	logs.Close()
	errs.Close()
	took := time.Since(start).Round(time.Second)
	if verdict != nil {
		r.cfg.Logf("%s %s @ %.12s: FAIL in %s: %v", taskType, in.Repo, in.Checkin, took, verdict)
		r.resolve(ctx, t.TaskID, map[string]any{"action": "fail", "error": verdict.Error()})
		return
	}
	r.cfg.Logf("%s %s @ %.12s: ok in %s", taskType, in.Repo, in.Checkin, took)
	r.resolve(ctx, t.TaskID, map[string]any{"action": "complete", "output": "ok"})
}

func (r *Runner) execute(ctx context.Context, taskType, checkin string, logs, errs io.Writer) error {
	if checkin == "" {
		return errors.New("the task names no check-in")
	}
	dir := filepath.Join(r.cfg.WorkDir, "ws", safeName(taskType))
	mu, _ := r.locks.LoadOrStore(dir, &sync.Mutex{})
	mu.(*sync.Mutex).Lock()
	defer mu.(*sync.Mutex).Unlock()
	fmt.Fprintf(logs, "$ checkout %.12s\n", checkin)
	if err := r.cfg.Checkout.Checkout(ctx, dir, checkin); err != nil {
		return fmt.Errorf("checkout: %w", err)
	}
	raw, err := os.ReadFile(filepath.Join(dir, ManifestPath))
	if err != nil {
		return fmt.Errorf("%s at %.12s: %v", ManifestPath, checkin, err)
	}
	m, err := ParseManifest(raw)
	if err != nil {
		return err
	}
	task, ok := m.Tasks[taskType]
	if !ok {
		return fmt.Errorf("%s at %.12s names no task %q (it names %s)", ManifestPath, checkin, taskType, strings.Join(m.Names(), ", "))
	}
	fmt.Fprintf(logs, "$ %s\n", strings.Join(task.Argv, " "))
	cmd := exec.CommandContext(ctx, task.Argv[0], task.Argv[1:]...)
	cmd.Dir = dir
	var caught bytes.Buffer
	cmd.Stdout = logs
	if task.FailOnOutput {
		cmd.Stdout = io.MultiWriter(logs, &caught)
	}
	cmd.Stderr = errs
	if err := cmd.Run(); err != nil {
		return err
	}
	if task.FailOnOutput && strings.TrimSpace(caught.String()) != "" {
		n := len(strings.Fields(caught.String()))
		fmt.Fprintf(errs, "%d line(s) of output from a task that must print nothing\n", n)
		return fmt.Errorf("%d line(s) of output from a task that must print nothing", n)
	}
	return nil
}

// heartbeat extends the task's claim while it runs.
func (r *Runner) heartbeat(ctx context.Context, taskID string) {
	tick := time.NewTicker(r.cfg.Heartbeat)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
			r.resolve(ctx, taskID, map[string]any{"action": "heartbeat"})
		}
	}
}

func (r *Runner) resolve(ctx context.Context, taskID string, body map[string]any) {
	b, _ := json.Marshal(body)
	resp, err := r.post(ctx, "/v1/tasks/"+taskID+"/resolve", b)
	if err != nil {
		if ctx.Err() == nil {
			r.cfg.Logf("resolve %s (%v): %v", taskID, body["action"], err)
		}
		return
	}
	resp.Body.Close()
}

func (r *Runner) post(ctx context.Context, path string, body []byte) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, "POST", strings.TrimSuffix(r.cfg.Bridge, "/")+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+r.cfg.Token)
	req.Header.Set("Content-Type", "application/json")
	resp, err := r.cfg.Client.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 300 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		resp.Body.Close()
		return nil, fmt.Errorf("%s: %d %s", path, resp.StatusCode, strings.TrimSpace(string(msg)))
	}
	return resp, nil
}

// logSink batches one stream of a task's output into the run's log: a
// flush every half second or 32 KB, whichever comes first, so a quiet
// step still shows progress and a loud one does not make a request per
// line. The engine names its two streams "out" and "err"; anything else
// is refused (measured: "stdout" was, so the first live run showed no
// output at all).
type logSink struct {
	r      *Runner
	stream string
	ctx    context.Context
	taskID string
	mu     sync.Mutex
	buf    bytes.Buffer
	warned bool
	done   chan struct{}
	wg     sync.WaitGroup
}

func newLogSink(ctx context.Context, r *Runner, taskID, stream string) *logSink {
	s := &logSink{r: r, ctx: ctx, taskID: taskID, stream: stream, done: make(chan struct{})}
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		tick := time.NewTicker(500 * time.Millisecond)
		defer tick.Stop()
		for {
			select {
			case <-s.done:
				s.flush()
				return
			case <-tick.C:
				s.flush()
			}
		}
	}()
	return s
}

func (s *logSink) Write(p []byte) (int, error) {
	s.mu.Lock()
	s.buf.Write(p)
	big := s.buf.Len() >= 32<<10
	s.mu.Unlock()
	if big {
		s.flush()
	}
	return len(p), nil
}

func (s *logSink) flush() {
	s.mu.Lock()
	if s.buf.Len() == 0 {
		s.mu.Unlock()
		return
	}
	data := append([]byte(nil), s.buf.Bytes()...)
	s.buf.Reset()
	s.mu.Unlock()
	body, _ := json.Marshal(map[string]any{"chunks": []map[string]string{{"stream": s.stream, "data": base64.StdEncoding.EncodeToString(data)}}})
	resp, err := s.r.post(s.ctx, "/v1/tasks/"+s.taskID+"/logs", body)
	if err != nil {
		// A lost batch does not fail the task, but it must not be
		// silent: a run page that shows no output while the runner
		// printed plenty is the confident-nothing shape. Said once per
		// task, so a dead log lane does not flood the runner's output.
		s.mu.Lock()
		first := !s.warned
		s.warned = true
		s.mu.Unlock()
		if first {
			s.r.cfg.Logf("logs for %s could not be sent (the run page will show none): %v", s.taskID, err)
		}
		return
	}
	resp.Body.Close()
}

func (s *logSink) Close() {
	close(s.done)
	s.wg.Wait()
}

func safeName(s string) string {
	if s == "" {
		return "_"
	}
	return strings.Map(func(c rune) rune {
		switch {
		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9', c == '-', c == '_', c == '.':
			return c
		}
		return '-'
	}, s)
}