package-updater/internal/agent/updater/runner.go

package updater

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sort"
	"sync"
	"time"

	"github.com/shcizo/package-updater/internal/protocol"
)

// maxProgress caps the output buffered between two TakeProgress calls, so a
// 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 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 }

// Finished is a job's outcome, kept until the master acknowledges it.
type Finished struct {
	ID     string
	Result protocol.JobResult
}

// tailBuffer keeps the last max bytes written to it.
type tailBuffer struct {
	max       int
	buf       []byte
	truncated bool
}

func (b *tailBuffer) Write(p []byte) {
	b.buf = append(b.buf, p...)
	if over := len(b.buf) - b.max; over > 0 {
		b.buf = append(b.buf[:0], b.buf[over:]...)
		b.truncated = true
	}
}

func (b *tailBuffer) String() string  { return string(b.buf) }
func (b *tailBuffer) Truncated() bool { return b.truncated }
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
	unsent   tailBuffer // since the last TakeProgress, capped at maxProgress
	reported bool       // a Progress has been taken since the job started
	finished time.Time  // orders Pending
}

// Runner executes jobs one at a time in the order they arrive and keeps each
// job's result until Ack. Two docker compose calls against the same stack
// would race, so there is deliberately exactly one worker.
//
// Runner satisfies conn.Jobs.
type Runner struct {
	exec    Executor
	timeout time.Duration
	now     func() time.Time

	mu    sync.Mutex // guards jobs and queue
	jobs  map[string]*jobState
	queue []string

	wake    chan struct{} // wakes the worker
	changed chan struct{} // tells the connection there is something to send
}

// NewRunner returns a Runner that gives every job at most timeout. Call Run
// to start the worker.
func NewRunner(exec Executor, timeout time.Duration) *Runner {
	return &Runner{
		exec:    exec,
		timeout: timeout,
		now:     time.Now,
		jobs:    make(map[string]*jobState),
		wake:    make(chan struct{}, 1),
		changed: make(chan struct{}, 1),
	}
}

// signal does a non-blocking send so neither side ever waits on the other.
func signal(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

// Submit queues a job. A job whose ID the runner already holds (queued,
// running, or finished and not yet acknowledged) is ignored.
func (r *Runner) Submit(id string, start protocol.JobStart) {
	r.mu.Lock()
	if _, ok := r.jobs[id]; ok {
		r.mu.Unlock()
		return
	}
	r.jobs[id] = &jobState{
		start:  start,
		output: tailBuffer{max: protocol.MaxJobOutput},
		unsent: tailBuffer{max: maxProgress},
	}
	r.queue = append(r.queue, id)
	r.mu.Unlock()
	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) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if j, ok := r.jobs[id]; ok && j.result != nil {
		delete(r.jobs, id)
	}
}

// Known returns the sorted IDs of every job the runner holds.
func (r *Runner) Known() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	ids := make([]string, 0, len(r.jobs))
	for id := range r.jobs {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	return ids
}

// TakeProgress returns the output running jobs produced since the previous
// call. A job's first entry (possibly with empty Output) tells the master it
// is running.
func (r *Runner) TakeProgress() []Progress {
	r.mu.Lock()
	defer r.mu.Unlock()
	var out []Progress
	for id, j := range r.jobs {
		if !j.running {
			continue
		}
		s := j.unsent.String()
		if j.reported && s == "" {
			continue
		}
		j.reported = true
		j.unsent.Reset()
		out = append(out, Progress{ID: id, Output: s})
	}
	sort.Slice(out, func(a, b int) bool { return out[a].ID < out[b].ID })
	return out
}

// Pending returns the results not yet acknowledged, oldest first.
func (r *Runner) Pending() []Finished {
	r.mu.Lock()
	defer r.mu.Unlock()
	var out []Finished
	var at []time.Time
	for id, j := range r.jobs {
		if j.result != nil {
			out = append(out, Finished{ID: id, Result: *j.result})
			at = append(at, j.finished)
		}
	}
	idx := make([]int, len(out))
	for i := range idx {
		idx[i] = i
	}
	sort.Slice(idx, func(a, b int) bool {
		if !at[idx[a]].Equal(at[idx[b]]) {
			return at[idx[a]].Before(at[idx[b]])
		}
		return out[idx[a]].ID < out[idx[b]].ID
	})
	sorted := make([]Finished, len(out))
	for i, k := range idx {
		sorted[i] = out[k]
	}
	return sorted
}

// Changed returns a channel that receives when there is new progress or a
// new result to send.
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. 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, ok := r.next()
		if !ok {
			select {
			case <-ctx.Done():
				return
			case <-r.wake:
				continue
			}
		}
		if ctx.Err() != nil {
			return
		}
		r.runOne(ctx, id)
	}
}

// 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 "", false
	}
	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, true
}

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)
	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:
		res = protocol.JobResult{Status: protocol.ResultFailed, Error: fmt.Sprintf("timed out after %s", r.timeout), StartedAt: started, FinishedAt: res.FinishedAt}
	case err != nil:
		res = protocol.JobResult{Status: protocol.ResultFailed, Error: err.Error(), StartedAt: started, FinishedAt: res.FinishedAt}
	}

	r.mu.Lock()
	res.Output, res.OutputTruncated = j.output.String(), j.output.Truncated()
	j.running, j.result, j.finished = false, &res, res.FinishedAt
	j.unsent.Reset()
	r.mu.Unlock()
	signal(r.changed)
}

// jobWriter feeds an executor's output into the job's buffers.
type jobWriter struct {
	r  *Runner
	id string
}

func (w *jobWriter) Write(p []byte) (int, error) {
	w.r.mu.Lock()
	if j, ok := w.r.jobs[w.id]; ok {
		j.output.Write(p)
		j.unsent.Write(p)
	}
	w.r.mu.Unlock()
	signal(w.r.changed)
	return len(p), nil
}