mesh-sdk/go/buildstep/pull.go

package buildstep

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"git.novox.be/novox/mesh-sdk/go/delivery"
)

// **The pull loop the build modules share** (novox/hq ADR 0525 §4-§6, ADR 0536 §4; issue 1887). A build module on a node
// whose build capacity is above zero pulls the asks of its language's seat: from the delivery worker first, then check,
// then recheck, holding at most its capacity at once. Each ask runs as a step under the ask's own id, so its end is the
// key <node>.<ask id> in the module's steps state, which the asker watches; the ask is acknowledged when that end is
// written, kept alive while it runs by an in-progress signal, and redelivered by the stream to another node if this one
// stops. No part keeps a timer for an ask: the pull's own wait and the stream's acknowledgement wait are the bus's.

// Pulled is one ask a source handed over, with what decides its fate on the bus.
type Pulled struct {
	// Kind is the worker it came from, filled by the puller.
	Kind string
	Data []byte
	// Delivered is how many times the bus has handed it over, this time included.
	Delivered uint64
	// Ack ends it: it is never delivered again. InProgress tells the bus it is still being worked, restarting its
	// acknowledgement wait. Nak hands it back to be delivered again at once (to this node or another). Term drops it.
	Ack, InProgress, Nak, Term func() error
}

// AskSource is where a puller takes asks from: the seat's workers on the bus (JetStreamAsks) or a test double
// (FakeAsks).
type AskSource interface {
	// Next takes one ask of kind: with wait 0 only one already waiting, else waiting for one at most wait. It answers
	// nil and no error when none came.
	Next(ctx context.Context, kind string, wait time.Duration) (*Pulled, error)
}

// Defaults of a Puller.
var (
	// DefaultPullWait is how long one pull of the delivery worker waits for an ask while every worker is empty: the
	// pull's own wait (as process.Consume's), after which the check and recheck workers are looked at again.
	DefaultPullWait = 5 * time.Second
	// DefaultHeartbeat is how often a running ask tells the bus it is still being worked; under the seat worker's
	// acknowledgement wait of 60 s (mesh-controller internal/broker HolderConsumerFor).
	DefaultHeartbeat = 20 * time.Second
)

// Puller pulls a build seat's asks and runs them.
type Puller struct {
	Source   AskSource
	Language string // the seat's language: delivery.AskGo, AskNode or AskImage
	Node     string // the node this holder runs on, the first part of its step keys
	Put      Put    // writes the module's steps state
	// Capacity is how many asks the node runs at once; read before every pull, so a change takes effect at the next
	// one. 0 pulls nothing.
	Capacity func() int
	// Run runs one ask's job; a *Refusal ends its step refused. Its result is the step's result.
	Run func(ctx context.Context, ask delivery.StepAsk) (any, error)
	// PullWait and Heartbeat default to DefaultPullWait and DefaultHeartbeat.
	PullWait  time.Duration
	Heartbeat time.Duration
	Logf      func(format string, args ...any)
	// Load reads one key of the module's steps state (nil and no error when it is absent). When set, Pull reads the
	// node's pause key (delivery.PauseKey) before it pulls, and Pause and Resume write it with Put, so a pause outlives
	// a restart.
	Load func(key string) ([]byte, error)

	mu      sync.Mutex
	inHand  int
	running map[string]RunningAsk
	paused  bool
	wake    chan struct{}
	once    sync.Once
}

func (p *Puller) init() {
	p.once.Do(func() { p.wake = make(chan struct{}, 1) })
}

func (p *Puller) logf(format string, args ...any) {
	if p.Logf != nil {
		p.Logf(format, args...)
	}
}

// Wake makes a waiting puller look again: after a capacity change, a resume, or a step's end.
func (p *Puller) Wake() {
	p.init()
	select {
	case p.wake <- struct{}{}:
	default:
	}
}

// Pause stops the puller taking asks; the asks it holds run to their end. by names who paused it. With Load set the
// pause is written on the node's pause key first, and an error there leaves the puller as it was.
func (p *Puller) Pause(by string) error { return p.setPaused(true, by) }

// Resume lets the puller take asks again, written as Pause's is.
func (p *Puller) Resume(by string) error { return p.setPaused(false, by) }

func (p *Puller) setPaused(paused bool, by string) error {
	if p.Load != nil {
		key, err := delivery.PauseKey(p.Node)
		if err != nil {
			return err
		}
		raw, err := delivery.EncodePause(delivery.Pause{Paused: paused, By: by, At: now()})
		if err != nil {
			return err
		}
		if _, err := p.Put(key, json.RawMessage(raw)); err != nil {
			return fmt.Errorf("buildstep: writing the pause of %s: %w", p.Node, err)
		}
	}
	p.mu.Lock()
	p.paused = paused
	p.mu.Unlock()
	p.Wake()
	return nil
}

// restore reads the node's pause key, when Load is set.
func (p *Puller) restore() error {
	if p.Load == nil {
		return nil
	}
	key, err := delivery.PauseKey(p.Node)
	if err != nil {
		return err
	}
	raw, err := p.Load(key)
	if err != nil {
		return fmt.Errorf("buildstep: reading the pause of %s: %w", p.Node, err)
	}
	if raw == nil {
		return nil
	}
	v, err := delivery.DecodePause(raw)
	if err != nil {
		return fmt.Errorf("buildstep: the pause of %s: %w", p.Node, err)
	}
	p.mu.Lock()
	p.paused = v.Paused
	p.mu.Unlock()
	return nil
}

// Paused says whether the puller takes no asks.
func (p *Puller) Paused() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.paused
}

// InHand is how many asks run now.
func (p *Puller) InHand() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.inHand
}

// RunningAsk is an ask a puller runs now, as its seat's current verb says it.
type RunningAsk struct {
	ID      string    `json:"id"`
	Kind    string    `json:"kind"`
	Module  string    `json:"module,omitempty"`
	Started time.Time `json:"started"`
}

// Running are the asks the puller runs now, oldest first.
func (p *Puller) Running() []RunningAsk {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]RunningAsk, 0, len(p.running))
	for _, r := range p.running {
		out = append(out, r)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Started.Before(out[j].Started) })
	return out
}

// room says whether the puller may take one more ask now.
func (p *Puller) room() bool {
	capacity := 0
	if p.Capacity != nil {
		capacity = p.Capacity()
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return !p.paused && p.inHand < capacity
}

// Pull takes and runs asks until ctx ends. An ask still running then is neither ended nor acknowledged: the stream
// delivers it again, to another node or to this one restarted.
func (p *Puller) Pull(ctx context.Context) error {
	p.init()
	if p.Source == nil || p.Put == nil || p.Run == nil {
		return errors.New("buildstep: a puller needs its source, its steps' put and its run")
	}
	if err := p.restore(); err != nil {
		return err // pulling on a pause that cannot be read could build on a node paused on purpose
	}
	wait := p.PullWait
	if wait <= 0 {
		wait = DefaultPullWait
	}
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		if !p.room() {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-p.wake:
			}
			continue
		}
		msg, err := p.next(ctx, wait)
		if err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			p.logf("the %s asks cannot be pulled: %v", p.Language, err)
			select { // the pull's own wait before the next try, as an empty pull's
			case <-ctx.Done():
				return ctx.Err()
			case <-p.wake:
			case <-time.After(wait):
			}
			continue
		}
		if msg == nil {
			continue
		}
		p.take(ctx, msg)
	}
}

// next is one ask from the workers in their order: one already waiting on delivery, check or recheck, and when none
// waits, a pull of the delivery worker that waits for one.
func (p *Puller) next(ctx context.Context, wait time.Duration) (*Pulled, error) {
	for _, kind := range delivery.AskKinds {
		msg, err := p.Source.Next(ctx, kind, 0)
		if err != nil || msg != nil {
			return withKind(msg, kind), err
		}
	}
	msg, err := p.Source.Next(ctx, delivery.AskDelivery, wait)
	return withKind(msg, delivery.AskDelivery), err
}

func withKind(msg *Pulled, kind string) *Pulled {
	if msg != nil {
		msg.Kind = kind
	}
	return msg
}

// take starts one ask: it decodes it, refuses one that is not of this seat, and runs it as a step under its own id.
func (p *Puller) take(ctx context.Context, msg *Pulled) {
	kind := msg.Kind
	ask, err := delivery.DecodeStepAsk(msg.Data)
	if err == nil && (ask.Language != p.Language || (kind != "" && ask.Kind != kind)) {
		err = fmt.Errorf("the ask %s is a %s ask of %s, pulled from the %s %s worker", ask.ID, ask.Kind, ask.Language, p.Language, kind)
	}
	if err != nil {
		// An ask that is not one runs nowhere: delivering it again would only fail again.
		p.logf("an ask dropped: %v", err)
		_ = msg.Term()
		return
	}
	step, key, err := begin(p.Put, p.Node, ask.Kind+" ask", ask.ID)
	if err != nil {
		p.logf("the ask %s cannot start on %s: %v", ask.ID, p.Node, err)
		_ = msg.Nak()
		return
	}
	module := ""
	if ask.Build != nil {
		module = ask.Build.Module
	}
	p.mu.Lock()
	p.inHand++
	if p.running == nil {
		p.running = map[string]RunningAsk{}
	}
	p.running[ask.ID] = RunningAsk{ID: ask.ID, Kind: ask.Kind, Module: module, Started: step.Started}
	p.mu.Unlock()
	go p.run(ctx, msg, ask, key, step)
}

func (p *Puller) run(ctx context.Context, msg *Pulled, ask delivery.StepAsk, key string, step delivery.Step) {
	defer func() {
		p.mu.Lock()
		p.inHand--
		delete(p.running, ask.ID)
		p.mu.Unlock()
		p.Wake()
	}()
	beat := p.Heartbeat
	if beat <= 0 {
		beat = DefaultHeartbeat
	}
	done := make(chan struct{})
	go func() {
		t := time.NewTicker(beat)
		defer t.Stop()
		for {
			select {
			case <-done:
				return
			case <-t.C:
				if msg.InProgress != nil {
					_ = msg.InProgress()
				}
			}
		}
	}()
	result, err := protect(ctx, func(ctx context.Context) (any, error) { return p.Run(ctx, ask) })
	close(done)
	if ctx.Err() != nil {
		// The process stops: the ask is neither ended nor acknowledged here, so the stream gives it to a node that runs.
		return
	}
	finish(p.Put, key, step, result, err)
	if aerr := msg.Ack(); aerr != nil {
		p.logf("the ask %s ended on %s and its acknowledgement failed: %v", ask.ID, p.Node, aerr)
	}
}