1// SPDX-License-Identifier: MIT
2
3package host
4
5import (
6 "context"
7 "sort"
8 "sync"
9 "time"
10
11 "connectrpc.com/connect"
12 "github.com/bnema/zerowrap"
13 "google.golang.org/protobuf/types/known/timestamppb"
14
15 ermite "bnema.dev/ermite-api/gen/ermite"
16)
17
18// Limits of AppendLog (runner.proto): at most 1,000 lines and 64 KiB of text
19// per request, all in the same step.
20const (
21 maxBatchLines = 1000
22 maxBatchBytes = 64 * 1024
23 // maxQueuedBytes is the log kept in memory while the server cannot be
24 // reached; beyond it the executor is slowed down (the stream is not read)
25 // rather than the host's memory filled.
26 maxQueuedBytes = 16 << 20
27)
28
29type logLine struct {
30 step int32
31 unixNano int64
32 text string
33 continued bool
34}
35
36// event is something to send, in the order it happened: a log line, or a
37// snapshot of a step (UpdateJob).
38type event struct {
39 line *logLine
40 step *ermite.StepUpdate
41 // anns are the annotations that go with a step snapshot, by step number:
42 // the whole list of each step they name (replace semantics).
43 anns map[int32][]*ermite.Annotation
44 // summary is a step summary for UploadSummary.
45 summary *summary
46}
47
48// summary is the Markdown of one step.
49type summary struct {
50 step int32
51 markdown string
52}
53
54// outbox sends logs and step updates of one job to ermite, in order, one call
55// at a time. Log lines are numbered from 0 across the whole job; the number of
56// the next line is the offset of the next AppendLog, so that a retry after a
57// lost response sends the same lines at the same offset (the server skips what
58// it already holds).
59type outbox struct {
60 client callClient
61 jobID string
62 t Timing
63 // onEnded is called once when the server says the job is over (JobEnded)
64 // or the assignment is gone; the outbox then drops everything.
65 onEnded func(error)
66 log zerowrap.Logger
67
68 mu sync.Mutex
69 cond *sync.Cond
70 events []event
71 // queued is the memory the outbox holds for the server, in bytes: log text
72 // and summaries in events, and annotations from the moment the host
73 // received them until the call that carries them is answered. One budget,
74 // maxQueuedBytes, for all of it.
75 queued int
76 closing bool
77 dead bool
78 // released is set by the hard stop of the job (releaseProducers): the
79 // producers that wait for room stop waiting and drop what they hold, and
80 // what is offered over the budget is dropped instead of waited for.
81 released bool
82 // stopSender ends the context the sender works with (start): kill calls
83 // it, so that a call or a retry delay in progress ends with the outbox and
84 // Close can wait for the sender after giving up.
85 stopSender context.CancelFunc
86 done chan struct{}
87
88 // pendingAnns are the annotations the server may not hold yet, by step: the
89 // last list queued for the step, until an UpdateJob that carried it
90 // succeeded. CompleteJob sends what is left.
91 pendingAnns map[int32][]*ermite.Annotation
92
93 offset int64 // lines the server holds, as far as we know
94 // offsetMaybeStale is set when a send was cut short by its context while
95 // the call may have been applied by the server with the response lost: the
96 // offset is then possibly behind the server's count.
97 offsetMaybeStale bool
98}
99
100// callClient is the part of the RunnerService client the job uses.
101type callClient interface {
102 UpdateJob(context.Context, *connect.Request[ermite.UpdateJobRequest]) (*connect.Response[ermite.UpdateJobResponse], error)
103 UploadSummary(context.Context, *connect.Request[ermite.UploadSummaryRequest]) (*connect.Response[ermite.UploadSummaryResponse], error)
104 AppendLog(context.Context, *connect.Request[ermite.AppendLogRequest]) (*connect.Response[ermite.AppendLogResponse], error)
105}
106
107func newOutbox(c callClient, jobID string, t Timing, onEnded func(error), log zerowrap.Logger) *outbox {
108 o := &outbox{client: c, jobID: jobID, t: t, onEnded: onEnded, log: log, done: make(chan struct{})}
109 o.cond = sync.NewCond(&o.mu)
110 return o
111}
112
113// start runs the sender until Close or until ctx ends.
114func (o *outbox) start(ctx context.Context) {
115 sctx, stop := context.WithCancel(ctx)
116 o.mu.Lock()
117 o.stopSender = stop
118 o.mu.Unlock()
119 go func() {
120 defer close(o.done)
121 defer stop()
122 o.loop(sctx)
123 }()
124 // A cancelled context must wake the loop and the producers.
125 go func() {
126 select {
127 case <-ctx.Done():
128 o.kill()
129 case <-o.done:
130 }
131 }()
132}
133
134func (o *outbox) kill() {
135 o.mu.Lock()
136 o.dead = true
137 o.events = nil
138 o.queued = 0
139 stop := o.stopSender
140 o.mu.Unlock()
141 if stop != nil {
142 stop()
143 }
144 o.cond.Broadcast()
145}
146
147// waitRoomLocked blocks while too much is queued. It returns false when the
148// item must be dropped: the outbox is dead, or the hard stop released the
149// producers and there is still no room. Closing lets the last items in.
150func (o *outbox) waitRoomLocked() bool {
151 for o.queued > maxQueuedBytes && !o.dead && !o.closing && !o.released {
152 o.cond.Wait()
153 }
154 if o.dead {
155 return false
156 }
157 return o.queued <= maxQueuedBytes || o.closing
158}
159
160// releaseProducers wakes the producers that wait for room and makes them (and
161// the ones that come) drop their items when there is none: called by the hard
162// stop of the job, because the sender may never make room (the server keeps
163// failing) and the producer is the reader of the stream of the sandbox. What
164// is queued is still sent by Close, within its budget.
165func (o *outbox) releaseProducers() {
166 o.mu.Lock()
167 o.released = true
168 o.mu.Unlock()
169 o.cond.Broadcast()
170}
171
172// Log queues a line of the job log. It blocks while too much is queued, until
173// the hard stop of the job releases it.
174func (o *outbox) Log(l logLine) {
175 o.mu.Lock()
176 if !o.waitRoomLocked() {
177 o.mu.Unlock()
178 return
179 }
180 o.events = append(o.events, event{line: &l})
181 o.queued += len(l.text) + 32
182 o.mu.Unlock()
183 o.cond.Broadcast()
184}
185
186// ReserveAnnotations counts the bytes of a list of annotations in the budget,
187// from the moment the host holds it. It returns false, and counts nothing,
188// when there is no room: the caller drops the list (annotations are not worth
189// stopping the reader of the stream for: it also carries the frames that
190// complete the steps and free the room).
191func (o *outbox) ReserveAnnotations(n int) bool {
192 o.mu.Lock()
193 defer o.mu.Unlock()
194 if o.dead || o.queued+n > maxQueuedBytes {
195 return false
196 }
197 o.queued += n
198 return true
199}
200
201// releaseAnnotations gives the bytes back once the call that carried the
202// lists was answered (or refused).
203func (o *outbox) releaseAnnotations(n int) {
204 o.mu.Lock()
205 if !o.dead {
206 o.queued -= n
207 }
208 o.mu.Unlock()
209 o.cond.Broadcast()
210}
211
212// Step queues a step snapshot, with the annotations that changed (may be nil).
213func (o *outbox) Step(s *ermite.StepUpdate, anns map[int32][]*ermite.Annotation) {
214 o.mu.Lock()
215 if o.dead {
216 o.mu.Unlock()
217 return
218 }
219 for n, l := range anns {
220 if o.pendingAnns == nil {
221 o.pendingAnns = map[int32][]*ermite.Annotation{}
222 }
223 o.pendingAnns[n] = l
224 }
225 o.events = append(o.events, event{step: s, anns: anns})
226 o.mu.Unlock()
227 o.cond.Broadcast()
228}
229
230// Summary queues the summary of a step. Its bytes count with the log text in
231// the bounded memory of the outbox: it blocks while too much is queued, until
232// the hard stop of the job releases it.
233func (o *outbox) Summary(step int32, markdown string) {
234 o.mu.Lock()
235 if !o.waitRoomLocked() {
236 o.mu.Unlock()
237 return
238 }
239 o.events = append(o.events, event{summary: &summary{step, markdown}})
240 o.queued += len(markdown) + 32
241 o.mu.Unlock()
242 o.cond.Broadcast()
243}
244
245// PendingAnnotations returns the annotations that no successful UpdateJob has
246// carried, for CompleteJob, in step order.
247func (o *outbox) PendingAnnotations() []*ermite.Annotation {
248 o.mu.Lock()
249 defer o.mu.Unlock()
250 return flattenAnnotations(o.pendingAnns)
251}
252
253func flattenAnnotations(m map[int32][]*ermite.Annotation) []*ermite.Annotation {
254 steps := make([]int32, 0, len(m))
255 for n := range m {
256 steps = append(steps, n)
257 }
258 sort.Slice(steps, func(a, b int) bool { return steps[a] < steps[b] })
259 var out []*ermite.Annotation
260 for _, n := range steps {
261 out = append(out, m[n]...)
262 }
263 return out
264}
265
266// Close sends what is queued, then stops the sender. It waits at most until
267// ctx ends.
268func (o *outbox) Close(ctx context.Context) {
269 o.mu.Lock()
270 o.closing = true
271 o.mu.Unlock()
272 o.cond.Broadcast()
273 select {
274 case <-o.done:
275 case <-ctx.Done():
276 o.kill()
277 <-o.done
278 }
279}
280
281// Lines returns how many lines the server holds (as far as we know).
282func (o *outbox) Lines() int64 {
283 o.mu.Lock()
284 defer o.mu.Unlock()
285 return o.offset
286}
287
288// next waits for a batch of events. It returns nil when the outbox is done.
289func (o *outbox) next(ctx context.Context) []event {
290 o.mu.Lock()
291 defer o.mu.Unlock()
292 for len(o.events) == 0 {
293 if o.dead || o.closing {
294 return nil
295 }
296 o.cond.Wait()
297 }
298 // Let a burst of lines fill the batch, unless it is full, a step event is
299 // at the front, or we are closing.
300 if o.events[0].line != nil && !o.closing && o.queued < maxBatchBytes && len(o.events) < maxBatchLines {
301 o.mu.Unlock()
302 sleep(ctx, o.t.LogFlush)
303 o.mu.Lock()
304 if o.dead {
305 return nil
306 }
307 }
308 if len(o.events) == 0 {
309 return nil
310 }
311 first := o.events[0]
312 if first.step != nil || first.summary != nil {
313 o.events = o.events[1:]
314 if first.summary != nil {
315 o.queued -= len(first.summary.markdown) + 32
316 o.cond.Broadcast() // producers waiting for room
317 }
318 return []event{first}
319 }
320 // A run of lines of the same step, within the limits.
321 n, bytes := 0, 0
322 step := first.line.step
323 for _, e := range o.events {
324 if e.line == nil || e.line.step != step || n == maxBatchLines || (n > 0 && bytes+len(e.line.text) > maxBatchBytes) {
325 break
326 }
327 n++
328 bytes += len(e.line.text)
329 }
330 batch := append([]event(nil), o.events[:n]...)
331 o.events = o.events[n:]
332 for _, e := range batch {
333 o.queued -= len(e.line.text) + 32
334 }
335 o.cond.Broadcast() // producers waiting for room
336 return batch
337}
338
339func (o *outbox) loop(ctx context.Context) {
340 for {
341 batch := o.next(ctx)
342 if batch == nil {
343 return
344 }
345 switch {
346 case batch[0].step != nil:
347 o.sendStep(ctx, batch[0].step, batch[0].anns)
348 case batch[0].summary != nil:
349 o.sendSummary(ctx, batch[0].summary)
350 default:
351 o.sendLines(ctx, batch)
352 }
353 o.mu.Lock()
354 dead := o.dead
355 o.mu.Unlock()
356 if dead {
357 return
358 }
359 }
360}
361
362// retry runs call until it succeeds, the server gives a verdict, or ctx ends.
363// It returns nil on success.
364func (o *outbox) retry(ctx context.Context, what string, call func(context.Context) error) error {
365 b := backoff{min: o.t.BackoffMin, max: o.t.BackoffMax}
366 for {
367 cctx, cancel := context.WithTimeout(ctx, o.t.CallTimeout)
368 err := call(cctx)
369 cancel()
370 if err == nil {
371 return nil
372 }
373 if ctx.Err() != nil {
374 return ctx.Err()
375 }
376 switch classify(err) {
377 case kindEnded, kindAbandoned:
378 o.onEnded(err)
379 o.kill()
380 return err
381 case kindPermanent:
382 return err
383 }
384 o.log.Warn().Err(err).Str("call", what).Msg("call failed, retrying")
385 if !sleep(ctx, b.next()) {
386 return ctx.Err()
387 }
388 }
389}
390
391func (o *outbox) sendStep(ctx context.Context, s *ermite.StepUpdate, anns map[int32][]*ermite.Annotation) {
392 req := &ermite.UpdateJobRequest{JobId: o.jobID, Steps: []*ermite.StepUpdate{s}, Annotations: flattenAnnotations(anns)}
393 err := o.retry(ctx, "UpdateJob", func(c context.Context) error {
394 _, err := o.client.UpdateJob(c, connect.NewRequest(req))
395 return err
396 })
397 if err == nil {
398 // The server holds these lists now: they leave the pending ones and
399 // the budget.
400 o.mu.Lock()
401 for n := range anns {
402 delete(o.pendingAnns, n)
403 }
404 o.mu.Unlock()
405 o.releaseAnnotations(annotationsBytes(anns))
406 return
407 }
408 if ctx.Err() == nil && classify(err) == kindPermanent {
409 o.log.Error().Err(err).Int32("step", s.Number).Msg("UpdateJob refused")
410 if len(anns) > 0 {
411 // The annotations may be what the server refused: the step update
412 // alone must not be lost with them.
413 o.mu.Lock()
414 for n := range anns {
415 delete(o.pendingAnns, n)
416 }
417 o.mu.Unlock()
418 o.releaseAnnotations(annotationsBytes(anns))
419 if err := o.retry(ctx, "UpdateJob", func(c context.Context) error {
420 _, err := o.client.UpdateJob(c, connect.NewRequest(&ermite.UpdateJobRequest{JobId: o.jobID, Steps: []*ermite.StepUpdate{s}}))
421 return err
422 }); err != nil && ctx.Err() == nil && classify(err) == kindPermanent {
423 o.log.Error().Err(err).Int32("step", s.Number).Msg("UpdateJob refused without annotations")
424 }
425 }
426 }
427}
428
429func (o *outbox) sendSummary(ctx context.Context, s *summary) {
430 err := o.retry(ctx, "UploadSummary", func(c context.Context) error {
431 _, err := o.client.UploadSummary(c, connect.NewRequest(&ermite.UploadSummaryRequest{JobId: o.jobID, StepNumber: s.step, Markdown: s.markdown}))
432 return err
433 })
434 if err != nil && ctx.Err() == nil && classify(err) == kindPermanent {
435 o.log.Error().Err(err).Int32("step", s.step).Msg("UploadSummary refused, the summary is dropped")
436 }
437}
438
439func (o *outbox) sendLines(ctx context.Context, batch []event) {
440 req := &ermite.AppendLogRequest{JobId: o.jobID, Step: batch[0].line.step}
441 for _, e := range batch {
442 req.Lines = append(req.Lines, &ermite.LogLine{Time: timestamppb.New(time.Unix(0, e.line.unixNano)), Text: e.line.text, Continued: e.line.continued})
443 }
444 // The offset is fixed before the first attempt: every retry sends the same
445 // lines at the same offset.
446 req.Offset = o.offset
447 var next int64
448 err := o.retry(ctx, "AppendLog", func(c context.Context) error {
449 resp, err := o.client.AppendLog(c, connect.NewRequest(req))
450 if err != nil {
451 return err
452 }
453 next = resp.Msg.NextOffset
454 return nil
455 })
456 switch {
457 case err == nil:
458 o.mu.Lock()
459 o.offsetMaybeStale = false
460 // The server says how many lines it holds: that is the truth. A
461 // difference means lines were stored that we did not count (a retry
462 // whose answer we lost) or dropped; continue from the server's count,
463 // so that the next offset is never a gap (out_of_range) and never a
464 // repeat of lines that would be skipped.
465 if want := req.Offset + int64(len(req.Lines)); next != want {
466 o.log.Warn().Int64("server_lines", next).Int64("expected", want).Msg("AppendLog: the server holds another number of lines, continuing from its count")
467 }
468 o.offset = next
469 o.mu.Unlock()
470 case ctx.Err() != nil:
471 // Cut short in the middle of a call or of its retries: the server may
472 // hold these lines without our having seen the answer.
473 o.mu.Lock()
474 o.offsetMaybeStale = true
475 o.mu.Unlock()
476 case classify(err) == kindPermanent:
477 // The server refused these lines; they are dropped and the offset stays
478 // where the server is.
479 o.log.Error().Err(err).Int("lines", len(req.Lines)).Int32("step", req.Step).Msg("AppendLog refused, the lines are dropped")
480 }
481}
482
483// AppendNow sends one line synchronously, after Close: the sender has stopped,
484// so nothing else uses the offset. It is for the last words of a job (an error
485// that explains why the result could not be reported).
486//
487// If Close gave up in the middle of a send, the offset may be behind the
488// server's: the line is sent anyway (the server answers with its own count, and
489// an offset that is too low makes it skip the line as a repeat or refuse it),
490// and the warning says why it may be missing.
491func (o *outbox) AppendNow(ctx context.Context, l logLine) {
492 o.mu.Lock()
493 stale := o.offsetMaybeStale
494 offset := o.offset
495 o.mu.Unlock()
496 if stale {
497 o.log.Warn().Int64("offset", offset).Msg("the log offset may be stale: Close was cut short in the middle of a send, so the last line may be skipped or refused by the server")
498 }
499 o.sendLines(ctx, []event{{line: &l}})
500}