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