mesh-tools/node-tools/internal/bus/seat_pull.go

package bus

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/nats-io/nats.go/jetstream"
)

// **A holder that pulls one piece of work when it has room** (novox/hq ADR 0525 §4, §6; issue 1887 A11). SeatTake
// hands a bundle every piece a worker holds, one at a time and prefetched; a build module must instead take one ask
// only when it has a free slot, from its workers in their order (delivery, then check, then recheck), and decide each
// ask's fate when its step ends. SeatPull is one fetch of at most one message from a worker the module's membership
// lists, never a prefetch; SeatSettle gives a pulled message its fate. A message the bundle never settles (its process
// died) is delivered again by the worker's own acknowledgement wait.

// MaxPullWait bounds how long one pull waits for a message: under the bundle's 30 s wait for its answer.
var MaxPullWait = 25 * time.Second

// Pulled is one message a module pulled from a seat's worker, named by an id the runtime gave it.
type Pulled struct {
	ID        string            `json:"id"`
	Worker    string            `json:"worker"`
	Subject   string            `json:"subject"`
	Body      json.RawMessage   `json:"body"`
	Headers   map[string]string `json:"headers,omitempty"`
	Delivered uint64            `json:"delivered"`
}

// pulls are the messages modules pulled and have not settled, by id; each kept with the module it was pulled for, so
// one module never settles another's.
type pulls struct {
	mu   sync.Mutex
	held map[string]pulledMsg
	cons map[string]jetstream.Consumer // by module and consumer
}

type pulledMsg struct {
	module string
	msg    jetstream.Msg
}

var seatPulls = &pulls{held: map[string]pulledMsg{}, cons: map[string]jetstream.Consumer{}}

func pullID() string {
	var b [8]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

// SeatPull fetches at most one message from consumer for module, waiting at most wait (none when 0): nil when none
// came. The consumer must be one the module's membership lists as its worker (novox/hq ADR 0259); the runtime binds to
// it by name and never makes one.
func (c *Conn) SeatPull(module, consumer string, wait time.Duration) (*Pulled, error) {
	t, err := c.trafficOf(module)
	if err != nil {
		return nil, err
	}
	var w *SeatWorker
	var names []string
	for i := range t.Workers {
		names = append(names, t.Workers[i].Consumer)
		if t.Workers[i].Consumer == consumer {
			w = &t.Workers[i]
		}
	}
	if w == nil {
		return nil, fmt.Errorf("%s pulls no work from %s: its membership lists %s (novox/hq ADR 0259)", module,
			consumer, strings.Join(names, ", "))
	}
	if wait > MaxPullWait {
		wait = MaxPullWait
	}
	bound, err := c.pullConsumer(module, *w)
	if err != nil {
		return nil, err
	}
	var batch jetstream.MessageBatch
	if wait <= 0 {
		batch, err = bound.FetchNoWait(1)
	} else {
		batch, err = bound.Fetch(1, jetstream.FetchMaxWait(wait))
	}
	if err != nil {
		return nil, fmt.Errorf("pulling %s's worker %s: %w", module, w.Consumer, err)
	}
	msg, ok := <-batch.Messages()
	if !ok {
		if err := batch.Error(); err != nil {
			return nil, fmt.Errorf("pulling %s's worker %s: %w", module, w.Consumer, err)
		}
		return nil, nil
	}
	if !json.Valid(msg.Data()) {
		_ = msg.Term()
		return nil, nil
	}
	headers := map[string]string{}
	for k := range msg.Headers() {
		headers[k] = msg.Headers().Get(k)
	}
	delivered := uint64(1)
	if meta, merr := msg.Metadata(); merr == nil && meta != nil {
		delivered = meta.NumDelivered
	}
	p := &Pulled{ID: pullID(), Worker: w.Consumer, Subject: msg.Subject(), Body: json.RawMessage(msg.Data()),
		Headers: headers, Delivered: delivered}
	seatPulls.mu.Lock()
	seatPulls.held[p.ID] = pulledMsg{module: module, msg: msg}
	seatPulls.mu.Unlock()
	return p, nil
}

func (c *Conn) pullConsumer(module string, w SeatWorker) (jetstream.Consumer, error) {
	key := module + " " + w.Consumer
	seatPulls.mu.Lock()
	bound, ok := seatPulls.cons[key]
	seatPulls.mu.Unlock()
	if ok {
		return bound, nil
	}
	js, err := jetstream.New(c.nc)
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	bound, err = js.Consumer(ctx, w.Stream, w.Consumer)
	if err != nil {
		return nil, fmt.Errorf("binding %s's worker %s on %s: %w", module, w.Consumer, w.Stream, err)
	}
	seatPulls.mu.Lock()
	seatPulls.cons[key] = bound
	seatPulls.mu.Unlock()
	return bound, nil
}

// The fates SeatSettle gives a pulled message.
const (
	SettleAck      = "ack"      // done: never delivered again
	SettleNak      = "nak"      // handed back, to be delivered again at once
	SettleTerm     = "term"     // dropped
	SettleProgress = "progress" // still being worked: its acknowledgement wait starts again
)

// SeatSettle gives the message module pulled under id its fate. A message pulled for another module, or settled
// already, is refused.
func (c *Conn) SeatSettle(module, id, fate string) error {
	seatPulls.mu.Lock()
	p, ok := seatPulls.held[id]
	if ok && p.module != module {
		ok = false
	}
	if ok && fate != SettleProgress {
		delete(seatPulls.held, id)
	}
	seatPulls.mu.Unlock()
	if !ok {
		return fmt.Errorf("%s holds no pulled message %s", module, id)
	}
	switch fate {
	case SettleAck:
		return p.msg.Ack()
	case SettleNak:
		return p.msg.Nak()
	case SettleTerm:
		return p.msg.Term()
	case SettleProgress:
		return p.msg.InProgress()
	}
	seatPulls.mu.Lock()
	seatPulls.held[id] = p // an unknown fate settles nothing
	seatPulls.mu.Unlock()
	return fmt.Errorf("a pulled message's fate is ack, nak, term or progress, not %q", fate)
}