mesh-delivery/modules/mesh-delivery-sdk/steprun/pull_bus.go

package steprun

import (
	"context"
	"fmt"
	"sync"
	"time"

	sdkstep "git.novox.be/novox/mesh-delivery/modules/mesh-delivery-sdk/buildstep"
	"github.com/nats-io/nats.go/jetstream"
)

// JetStreamAsks is a build seat's workers on the bus: the shared pull consumer of each kind on the seat's stream
// (sdkstep.AskStream, AskWorker), which the broker makes when the seat is declared; a holder only binds them.
type JetStreamAsks struct {
	JS       jetstream.JetStream
	Language string

	mu      sync.Mutex
	cons    map[string]jetstream.Consumer
	missing map[string]bool
}

func (s *JetStreamAsks) consumer(ctx context.Context, kind string, named bool) (jetstream.Consumer, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	cacheKey := kind
	if named {
		cacheKey = kind + "/named"
	}
	if c, ok := s.cons[cacheKey]; ok {
		return c, nil
	}
	stream, err := sdkstep.AskStream(s.Language)
	if err != nil {
		return nil, err
	}
	worker, err := sdkstep.AskWorker(s.Language, kind)
	if named {
		worker, err = sdkstep.AskNamedWorker(s.Language, kind)
	}
	if err != nil {
		return nil, err
	}
	c, err := s.JS.Consumer(ctx, stream, worker)
	if err != nil {
		return nil, fmt.Errorf("binding the %s worker %s on %s: %w", kind, worker, stream, err)
	}
	if s.cons == nil {
		s.cons = map[string]jetstream.Consumer{}
	}
	s.cons[cacheKey] = c
	return c, nil
}

// Next is AskSource's: main's worker first, then the queue's named worker once it is bound (design 47 G §544), with no
// wait when wait is 0, else waiting at most wait between the two. A named worker not bound is only main's alone.
func (s *JetStreamAsks) Next(ctx context.Context, kind string, wait time.Duration) (*Pulled, error) {
	var half time.Duration
	if wait > 0 {
		half = wait / 2
	}
	for i, step := range []struct {
		named bool
		wait  time.Duration
	}{{false, 0}, {true, 0}, {false, half}, {true, wait - half}} {
		if step.named && s.unbound(kind) {
			continue
		}
		if i >= 2 && wait <= 0 {
			break // the two waiting pulls are for a wait
		}
		c, err := s.consumer(ctx, kind, step.named)
		if err != nil {
			if step.named {
				s.markUnbound(kind)
				continue
			}
			return nil, err
		}
		if msg, err := s.fetch(ctx, c, step.wait); err != nil || msg != nil {
			return msg, err
		}
	}
	return nil, nil
}

func (s *JetStreamAsks) unbound(kind string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.missing[kind]
}

func (s *JetStreamAsks) markUnbound(kind string) {
	s.mu.Lock()
	if s.missing == nil {
		s.missing = map[string]bool{}
	}
	s.missing[kind] = true
	s.mu.Unlock()
}

func (s *JetStreamAsks) fetch(ctx context.Context, c jetstream.Consumer, wait time.Duration) (*Pulled, error) {
	var batch jetstream.MessageBatch
	var err error
	if wait <= 0 {
		batch, err = c.FetchNoWait(1)
	} else {
		batch, err = c.Fetch(1, jetstream.FetchMaxWait(wait))
	}
	if err != nil {
		return nil, err
	}
	select {
	case msg, ok := <-batch.Messages():
		if !ok {
			return nil, batch.Error()
		}
		var delivered uint64 = 1
		if meta, err := msg.Metadata(); err == nil && meta != nil {
			delivered = meta.NumDelivered
		}
		return &Pulled{Data: append([]byte(nil), msg.Data()...), Subject: msg.Subject(), Delivered: delivered,
			Ack: msg.Ack, InProgress: msg.InProgress, Nak: msg.Nak, Term: msg.Term,
			NakAfter: func(delay time.Duration) error { return msg.NakWithDelay(delay) }}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// DeclareAskWorkers makes language's seat stream and its three workers on a bus with no broker: a test's (the build
// modules' integration tests start one). On the mesh the broker declares them from the seat, as it declares every
// seat's stream (mesh-controller internal/broker SeatStreams); a holder never does.
func DeclareAskWorkers(ctx context.Context, js jetstream.JetStream, language string, ackWait time.Duration) error {
	stream, err := sdkstep.AskStream(language)
	if err != nil {
		return err
	}
	seat, _ := sdkstep.BuildSeat(language)
	if _, err := js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{Name: stream, Subjects: []string{"mesh.seat." + seat + ".accept.>"},
		Retention: jetstream.WorkQueuePolicy}); err != nil {
		return fmt.Errorf("the stream %s: %w", stream, err)
	}
	for _, kind := range sdkstep.AskKinds {
		worker, _ := sdkstep.AskWorker(language, kind)
		subject, _ := sdkstep.AskSubject(language, kind)
		// Two workers whose filters do not overlap: main's subject and the asker's, `<main's>.<asker>` (design 47 G §544).
		if _, err := js.CreateOrUpdateConsumer(ctx, stream, jetstream.ConsumerConfig{Durable: worker, FilterSubjects: []string{subject},
			AckPolicy: jetstream.AckExplicitPolicy, AckWait: ackWait, MaxDeliver: 5}); err != nil {
			return fmt.Errorf("the worker %s: %w", worker, err)
		}
		named, _ := sdkstep.AskNamedWorker(language, kind)
		if _, err := js.CreateOrUpdateConsumer(ctx, stream, jetstream.ConsumerConfig{Durable: named, FilterSubjects: []string{subject + ".*"},
			AckPolicy: jetstream.AckExplicitPolicy, AckWait: ackWait, MaxDeliver: 5}); err != nil {
			return fmt.Errorf("the worker %s: %w", named, err)
		}
	}
	return nil
}