1// SPDX-License-Identifier: MIT
2
3package host
4
5import (
6 "context"
7 "crypto/rand"
8 "encoding/base64"
9 "errors"
10 "fmt"
11 "runtime"
12 "sync"
13 "sync/atomic"
14 "time"
15
16 "connectrpc.com/connect"
17 "github.com/bnema/zerowrap"
18
19 ermite "bnema.dev/ermite-api/gen/ermite"
20 "bnema.dev/ermite-api/gen/ermite/ermiteconnect"
21
22 "bnema.dev/ermite-runner/internal/execstream/execpb"
23 "bnema.dev/ermite-runner/internal/sandbox"
24)
25
26// Runner fetches jobs from ermite and runs each one in its own sandbox.
27type Runner struct {
28 Cfg Config
29 Driver sandbox.Driver
30 // Poll is the client that carries the runner token (FetchJob).
31 Poll ermiteconnect.RunnerServiceClient
32 // Job builds the client of a job from its job token.
33 Job func(jobToken string) ermiteconnect.RunnerServiceClient
34 Timing Timing
35 // Features lists the features of the server (GetVersion). Annotations and
36 // step summaries are sent only when it says "ci.annotations" and
37 // "ci.summaries". Nil means neither is sent.
38 Features ServerFeatures
39
40 // now returns the current time; tests replace it.
41 now func() time.Time
42
43 leaked atomic.Bool
44 active sync.WaitGroup
45}
46
47func (r *Runner) clock() time.Time {
48 if r.now != nil {
49 return r.now()
50 }
51 return time.Now()
52}
53
54// spec is the sandbox request of a job. It names the job and leaves CPUs,
55// memory and disk at zero: the driver fills them from its configuration
56// (sandbox.Config.Resolve in Start), that is from ERMITE_RUNNER_JOB_CPUS,
57// _JOB_MEMORY and _JOB_DISK, so that the limits are one setting of the runner
58// and not something a job or this package can change.
59func (r *Runner) spec(j *ermite.Job) sandbox.JobSpec {
60 return sandbox.JobSpec{JobID: j.JobId}
61}
62
63// markLeaked records that a sandbox could not be cleaned up: no new job is
64// taken until the running ones are over and Driver.Reap has succeeded.
65func (r *Runner) markLeaked() { r.leaked.Store(true) }
66
67// ErrRemoved is returned by Run when ermite no longer knows the runner token:
68// the runner was removed or replaced.
69var ErrRemoved = errors.New("ermite rejected the runner token: the runner was removed or replaced; register it again")
70
71// newPollID is 16 random bytes in base64url: 22 characters of the set the
72// protocol allows.
73func newPollID() string {
74 var b [16]byte
75 _, _ = rand.Read(b[:])
76 return base64.RawURLEncoding.EncodeToString(b[:])
77}
78
79// Run is the main loop. It returns nil after a graceful shutdown (ctx ended and
80// every running job reported), or an error when the runner cannot go on.
81//
82// Capacity: a poll is made only when a slot is free. A lost response is retried
83// with the same poll_id, so that the server hands back the same assignment.
84func (r *Runner) Run(ctx context.Context) error {
85 log := zerowrap.FromCtx(ctx)
86 t := r.Timing.withDefaults()
87 shutdown := make(chan struct{})
88 slots := make(chan struct{}, max(1, r.Cfg.Capacity))
89 defer func() {
90 // Graceful shutdown: running jobs are cancelled with runner_shutdown and
91 // get ShutdownGrace to stop; Run returns once they have reported.
92 close(shutdown)
93 r.active.Wait()
94 }()
95
96 if err := r.Driver.Reap(ctx); err != nil && !errors.Is(err, sandbox.ErrBusy) {
97 log.Error().Err(err).Msg("reaping what a previous run left")
98 }
99
100 var pollID string
101 bo := backoff{min: t.BackoffMin, max: t.BackoffMax}
102 for ctx.Err() == nil {
103 // Wait for a free slot, and for the end of a leak.
104 select {
105 case slots <- struct{}{}:
106 case <-ctx.Done():
107 return nil
108 }
109 if r.leaked.Load() {
110 <-slots
111 if !r.recover(ctx, log, slots, t.ReapRetry) {
112 return nil
113 }
114 continue
115 }
116 if pollID == "" {
117 pollID = newPollID()
118 }
119 pctx, cancel := context.WithTimeout(ctx, t.PollTimeout)
120 resp, err := r.Poll.FetchJob(pctx, connect.NewRequest(&ermite.FetchJobRequest{PollId: pollID}))
121 cancel()
122 if err != nil {
123 <-slots
124 if ctx.Err() != nil {
125 return nil
126 }
127 if connect.CodeOf(err) == connect.CodeUnauthenticated {
128 return ErrRemoved
129 }
130 if connect.CodeOf(err) == connect.CodeInvalidArgument {
131 pollID = "" // our id was refused: a new one
132 }
133 d := bo.next()
134 log.Warn().Err(err).Dur("retry_in", d).Msg("FetchJob failed")
135 if !sleep(ctx, d) {
136 return nil
137 }
138 continue
139 }
140 bo.cur = 0
141 pollID = "" // answered: the next poll is a new one
142 if resp.Msg.Job == nil {
143 <-slots
144 d := t.NoJobPause
145 if ra := resp.Msg.RetryAfter; ra != nil && ra.AsDuration() > d {
146 d = ra.AsDuration()
147 }
148 if !sleep(ctx, d) {
149 return nil
150 }
151 continue
152 }
153 r.startJob(ctx, log, resp.Msg.Job, t, shutdown, slots)
154 if ra := resp.Msg.RetryAfter; ra != nil && ra.AsDuration() > 0 {
155 if !sleep(ctx, ra.AsDuration()) {
156 return nil
157 }
158 }
159 }
160 return nil
161}
162
163// recover waits until no job runs, then reaps. It returns false when ctx ends.
164func (r *Runner) recover(ctx context.Context, log zerowrap.Logger, slots chan struct{}, retry time.Duration) bool {
165 log.Error().Msg("a sandbox could not be cleaned up: no new job until the running ones end and the leftovers are reaped")
166 for ctx.Err() == nil {
167 if len(slots) == 0 {
168 if err := r.Driver.Reap(ctx); err == nil {
169 r.leaked.Store(false)
170 log.Info().Msg("reaped: taking jobs again")
171 return true
172 } else if !errors.Is(err, sandbox.ErrBusy) {
173 log.Error().Err(err).Msg("reaping failed, retrying")
174 }
175 }
176 if !sleep(ctx, retry) {
177 return false
178 }
179 }
180 return false
181}
182
183func (r *Runner) startJob(ctx context.Context, log zerowrap.Logger, job *ermite.Job, t Timing, shutdown <-chan struct{}, slots chan struct{}) {
184 jctx, cancel := context.WithCancelCause(context.WithoutCancel(ctx))
185 jr := &jobRun{
186 r: r, job: job, client: r.Job(job.JobToken), t: t, ctx: jctx, cancel: cancel,
187 log: zerowrap.Logger{Logger: log.With().Str("job_id", job.JobId).Int64("run_id", job.RunId).Int32("attempt", job.Attempt).Logger()},
188 receivedAt: r.clock(), cancelCh: make(chan *execpb.Cancel, 1), steps: map[int32]*ermite.StepUpdate{},
189 annSeen: map[int32]bool{}, sumSeen: map[int32]bool{},
190 }
191 jr.log.Info().Str("job_key", job.JobKey).Msg("job assigned")
192 r.active.Add(1)
193 go func() {
194 defer r.active.Done()
195 defer func() { <-slots }()
196 jr.run(r, shutdown)
197 jr.log.Info().Msg("job finished")
198 }()
199}
200
201// ---- registration
202
203// Isolation maps the driver's isolation to the protocol's.
204func Isolation(i sandbox.Isolation) ermite.RunnerIsolation {
205 switch i {
206 case sandbox.IsolationVM:
207 return ermite.RunnerIsolation_RUNNER_ISOLATION_VM
208 case sandbox.IsolationContainerSysbox:
209 return ermite.RunnerIsolation_RUNNER_ISOLATION_CONTAINER_SYSBOX
210 case sandbox.IsolationContainerPrivileged:
211 return ermite.RunnerIsolation_RUNNER_ISOLATION_CONTAINER_PRIVILEGED
212 }
213 return ermite.RunnerIsolation_RUNNER_ISOLATION_UNSPECIFIED
214}
215
216// GoArch is the architecture in GitHub's names.
217func GoArch() string {
218 if runtime.GOARCH == "arm64" {
219 return "arm64"
220 }
221 return "x64"
222}
223
224// Register registers the runner with ermite using a token that has the
225// runner:register scope, and stores the runner token (mode 0600) in the state
226// directory. replace removes a runner of the same name first.
227func Register(ctx context.Context, cfg Config, client ermiteconnect.RunnerServiceClient, isolation sandbox.Isolation, replace bool) (State, error) {
228 iso := Isolation(isolation)
229 if iso == ermite.RunnerIsolation_RUNNER_ISOLATION_UNSPECIFIED {
230 return State{}, fmt.Errorf("unknown isolation %q", isolation)
231 }
232 rctx, cancel := context.WithTimeout(ctx, 30*time.Second)
233 defer cancel()
234 resp, err := client.RegisterRunner(rctx, connect.NewRequest(&ermite.RegisterRunnerRequest{
235 Name: cfg.Name, Labels: cfg.Labels, Version: cfg.Version, Isolation: iso, Arch: GoArch(), Os: "linux", Replace: replace,
236 }))
237 if err != nil {
238 return State{}, fmt.Errorf("register: %w", err)
239 }
240 st := State{URL: cfg.URL, RunnerID: resp.Msg.RunnerId, Name: cfg.Name, RunnerToken: resp.Msg.RunnerToken}
241 if err := SaveState(cfg.StateDir, st); err != nil {
242 return State{}, fmt.Errorf("the runner is registered but its token could not be stored: %w", err)
243 }
244 return st, nil
245}