feat: self-update of master and agents (#15) (#32)

internal/agent/updater/runner.go

@@ -16,12 +16,25 @@ import (
 // chatty job cannot grow a single job.progress message without bound.
 const maxProgress = 32 << 10
 
+// Job is a job as an Executor sees it: what the master asked for, plus the
+// ID and start time the runner gave it. A self-update records both in the
+// helper container, so the new agent can finish the job under the same ID.
+type Job struct {
+	protocol.JobStart
+	ID      string
+	Started time.Time
+}
+
 // Executor deploys one job and returns the "repo:tag@digest" it deployed.
 // Everything it writes to out is streamed to the master as job output.
 type Executor interface {
-	Execute(ctx context.Context, job protocol.JobStart, out io.Writer) (string, error)
+	Execute(ctx context.Context, job Job, out io.Writer) (string, error)
 }
 
+// ResumeFunc finishes a job that an earlier agent process began. It returns
+// the "repo:tag@digest" that was deployed, like Executor.Execute.
+type ResumeFunc func(ctx context.Context, out io.Writer) (string, error)
+
 // Progress is output a running job produced since the previous TakeProgress.
 type Progress struct{ ID, Output string }
 
@@ -52,6 +65,8 @@ func (b *tailBuffer) Reset()          { b.buf, b.truncated = nil, false }
 
 type jobState struct {
 	start    protocol.JobStart
+	resume   ResumeFunc // set for a job resumed after a restart; start is then empty
+	started  time.Time  // set with resume
 	running  bool
 	result   *protocol.JobResult
 	output   tailBuffer // everything, capped at protocol.MaxJobOutput
@@ -117,6 +132,28 @@ func (r *Runner) Submit(id string, start protocol.JobStart) {
 	signal(r.wake)
 }
 
+// Resume registers a job that an earlier agent process began and runs fn for
+// it, before every job that is already queued. The job is in Known at once,
+// so the agent's next hello does not report it lost. The result carries
+// started as its start time. A job whose ID the runner already holds is
+// ignored.
+func (r *Runner) Resume(id string, started time.Time, fn ResumeFunc) {
+	r.mu.Lock()
+	if _, ok := r.jobs[id]; ok {
+		r.mu.Unlock()
+		return
+	}
+	r.jobs[id] = &jobState{
+		resume:  fn,
+		started: started,
+		output:  tailBuffer{max: protocol.MaxJobOutput},
+		unsent:  tailBuffer{max: maxProgress},
+	}
+	r.queue = append([]string{id}, r.queue...)
+	r.mu.Unlock()
+	signal(r.wake)
+}
+
 // Ack forgets a finished job. Acks for unknown, queued or running jobs are
 // ignored.
 func (r *Runner) Ack(id string) {
@@ -196,10 +233,12 @@ func (r *Runner) Pending() []Finished {
 func (r *Runner) Changed() <-chan struct{} { return r.changed }
 
 // Run works through the queue until ctx ends. Jobs still queued then are
-// dropped; the master marks them lost on the agent's next hello.
+// dropped; the master marks them lost on the agent's next hello. A job that
+// ctx interrupts is dropped too, without a result: a self-updating agent is
+// stopped that way, and its successor reports the outcome.
 func (r *Runner) Run(ctx context.Context) {
 	for {
-		id, start, ok := r.next()
+		id, ok := r.next()
 		if !ok {
 			select {
 			case <-ctx.Done():
@@ -211,32 +250,62 @@ func (r *Runner) Run(ctx context.Context) {
 		if ctx.Err() != nil {
 			return
 		}
-		r.runOne(ctx, id, start)
+		r.runOne(ctx, id)
 	}
 }
 
-// next pops the next queued job and marks it running.
-func (r *Runner) next() (string, protocol.JobStart, bool) {
+// next pops the next job to run and marks it running. Jobs for the agent's
+// own service run last: a master and an agent on the same host share an
+// image, and the master's job must finish before the agent replaces itself.
+func (r *Runner) next() (string, bool) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if len(r.queue) == 0 {
-		return "", protocol.JobStart{}, false
+		return "", false
 	}
-	id := r.queue[0]
-	r.queue = r.queue[1:]
-	j := r.jobs[id]
-	j.running = true
+	pick := 0
+	for i, id := range r.queue {
+		if !r.jobs[id].start.Self {
+			pick = i
+			break
+		}
+	}
+	id := r.queue[pick]
+	r.queue = append(r.queue[:pick:pick], r.queue[pick+1:]...)
+	r.jobs[id].running = true
 	signal(r.changed)
-	return id, j.start, true
+	return id, true
 }
 
-func (r *Runner) runOne(ctx context.Context, id string, start protocol.JobStart) {
-	started := r.now()
+func (r *Runner) runOne(ctx context.Context, id string) {
+	r.mu.Lock()
+	j := r.jobs[id]
+	start, resume, started := j.start, j.resume, j.started
+	r.mu.Unlock()
+	if resume == nil {
+		started = r.now()
+	}
+
 	jobCtx, cancel := context.WithTimeout(ctx, r.timeout)
-	ref, err := r.exec.Execute(jobCtx, start, &jobWriter{r: r, id: id})
+	var ref string
+	var err error
+	if resume != nil {
+		ref, err = resume(jobCtx, &jobWriter{r: r, id: id})
+	} else {
+		ref, err = r.exec.Execute(jobCtx, Job{JobStart: start, ID: id, Started: started}, &jobWriter{r: r, id: id})
+	}
 	timedOut := errors.Is(jobCtx.Err(), context.DeadlineExceeded)
 	cancel()
 
+	if err != nil && ctx.Err() != nil && !timedOut {
+		// The agent is stopping, not the job failing. Nothing is recorded, so
+		// no result can pre-empt the one a successor agent may report.
+		r.mu.Lock()
+		delete(r.jobs, id)
+		r.mu.Unlock()
+		return
+	}
+
 	res := protocol.JobResult{Status: protocol.ResultSucceeded, Ref: ref, StartedAt: started, FinishedAt: r.now()}
 	switch {
 	case timedOut:
@@ -246,7 +315,6 @@ func (r *Runner) runOne(ctx context.Context, id string, start protocol.JobStart)
 	}
 
 	r.mu.Lock()
-	j := r.jobs[id]
 	res.Output, res.OutputTruncated = j.output.String(), j.output.Truncated()
 	j.running, j.result, j.finished = false, &res, res.FinishedAt
 	j.unsent.Reset()