gitea-dynamic-runner/internal/assignmentqueue/jetstream.go

// Package assignmentqueue implements the durable assignment handoff with JetStream.
package assignmentqueue

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

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

	"git.ddupan.top/panxiao81/gitea-dynamic-runner/internal/taskassignment"
)

type publishAPI interface {
	PublishMsg(context.Context, *nats.Msg, ...jetstream.PublishOpt) (*jetstream.PubAck, error)
}

// Publisher implements the scheduler dispatcher with one subject per backend.
type Publisher struct {
	JetStream   publishAPI
	SubjectBase string
}

func (p Publisher) Dispatch(ctx context.Context, assignment taskassignment.Assignment) error {
	if p.JetStream == nil {
		return errors.New("JetStream publisher is required")
	}
	body, err := taskassignment.Marshal(assignment)
	if err != nil {
		return err
	}
	base := strings.TrimSuffix(p.SubjectBase, ".")
	if base == "" {
		return errors.New("assignment subject base is required")
	}
	message := &nats.Msg{
		Subject: base + "." + string(assignment.Backend),
		Header:  nats.Header{jetstream.MsgIDHeader: []string{assignment.ID}},
		Data:    body,
	}
	if _, err := p.JetStream.PublishMsg(ctx, message); err != nil {
		return fmt.Errorf("publish assignment %s: %w", assignment.ID, err)
	}
	return nil
}

type Accepter interface {
	Accept(context.Context, taskassignment.Assignment) (bool, error)
}

type Claims interface {
	Offer(taskassignment.Assignment) (<-chan struct{}, error)
	WaitClaimed(context.Context, string) error
}

type Admission interface {
	Acquire(string) bool
	Release(string)
}

// Message is the subset of jetstream.Msg needed by one reconciliation.
type Message interface {
	Data() []byte
	DoubleAck(context.Context) error
	NakWithDelay(time.Duration) error
	TermWithReason(string) error
}

// Processor maps one delivery to one idempotent worker reconciliation.
type Processor struct {
	TrustDomain  string
	Accepter     Accepter
	Claims       Claims
	Admission    Admission
	RetryDelay   time.Duration
	ClaimTimeout time.Duration
}

func (p Processor) Process(ctx context.Context, message Message) error {
	if p.Accepter == nil || p.Claims == nil || p.Admission == nil {
		return errors.New("assignment accepter, claim registry, and backend admission pool are required")
	}
	assignment, err := taskassignment.Unmarshal(message.Data(), p.TrustDomain)
	if err != nil {
		return errors.Join(err, message.TermWithReason("invalid assignment"))
	}
	if _, err := p.Claims.Offer(assignment); err != nil {
		return errors.Join(err, message.TermWithReason("conflicting assignment"))
	}
	if !p.Admission.Acquire(assignment.ID) {
		delay := p.RetryDelay
		if delay <= 0 {
			delay = 2 * time.Second
		}
		return message.NakWithDelay(delay)
	}
	accepted, err := p.Accepter.Accept(ctx, assignment)
	if err != nil {
		p.Admission.Release(assignment.ID)
		delay := p.RetryDelay
		if delay <= 0 {
			delay = 15 * time.Second
		}
		return errors.Join(err, message.NakWithDelay(delay))
	}
	if accepted {
		timeout := p.ClaimTimeout
		if timeout <= 0 {
			timeout = 4 * time.Minute
		}
		claimContext, cancel := context.WithTimeout(ctx, timeout)
		err := p.Claims.WaitClaimed(claimContext, assignment.ID)
		cancel()
		if err != nil {
			delay := p.RetryDelay
			if delay <= 0 {
				delay = 2 * time.Second
			}
			return errors.Join(err, message.NakWithDelay(delay))
		}
		if err := message.DoubleAck(ctx); err != nil {
			return fmt.Errorf("ack assignment %s: %w", assignment.ID, err)
		}
		return nil
	}
	delay := p.RetryDelay
	if delay <= 0 {
		delay = 2 * time.Second
	}
	return message.NakWithDelay(delay)
}

type consumeAPI interface {
	Consume(jetstream.MessageHandler, ...jetstream.PullConsumeOpt) (jetstream.ConsumeContext, error)
}

// ConsumerComponent runs bounded reconciliation goroutines for one durable
// backend consumer. The goroutine set is operational state, not task storage.
type ConsumerComponent struct {
	Consumer  consumeAPI
	Processor Processor
	Capacity  int
	OnError   func(error)
}

type consumerManager interface {
	CreateOrUpdateConsumer(context.Context, string, jetstream.ConsumerConfig) (jetstream.Consumer, error)
}

// OpenConsumer creates the durable backend cursor. Capacity is enforced both
// server-side and by ConsumerComponent's local semaphore.
func OpenConsumer(ctx context.Context, manager consumerManager, stream, subjectBase string, backend taskassignment.Backend, capacity int) (jetstream.Consumer, error) {
	if manager == nil || stream == "" || strings.TrimSuffix(subjectBase, ".") == "" || capacity < 1 {
		return nil, errors.New("JetStream manager, stream, subject base, and positive capacity are required")
	}
	if backend != taskassignment.BackendPod && backend != taskassignment.BackendVM {
		return nil, fmt.Errorf("unsupported assignment backend %q", backend)
	}
	consumer, err := manager.CreateOrUpdateConsumer(ctx, stream, jetstream.ConsumerConfig{
		Name:          string(backend),
		Durable:       string(backend),
		FilterSubject: strings.TrimSuffix(subjectBase, ".") + "." + string(backend),
		AckPolicy:     jetstream.AckExplicitPolicy,
		AckWait:       5 * time.Minute,
		MaxAckPending: capacity,
		MaxDeliver:    20,
	})
	if err != nil {
		return nil, fmt.Errorf("open %s assignment consumer: %w", backend, err)
	}
	return consumer, nil
}

func (c ConsumerComponent) Run(ctx context.Context) error {
	if c.Consumer == nil || c.Capacity < 1 {
		return errors.New("JetStream consumer and positive capacity are required")
	}
	semaphore := make(chan struct{}, c.Capacity)
	var workers sync.WaitGroup
	consumeContext, err := c.Consumer.Consume(func(message jetstream.Msg) {
		select {
		case semaphore <- struct{}{}:
		case <-ctx.Done():
			return
		}
		workers.Add(1)
		go func() {
			defer workers.Done()
			defer func() { <-semaphore }()
			if err := c.Processor.Process(ctx, message); err != nil && !errors.Is(err, context.Canceled) && c.OnError != nil {
				c.OnError(err)
			}
		}()
	}, jetstream.PullMaxMessages(c.Capacity))
	if err != nil {
		return fmt.Errorf("start JetStream consumer: %w", err)
	}

	select {
	case <-ctx.Done():
		consumeContext.Stop()
		<-consumeContext.Closed()
		workers.Wait()
		return nil
	case <-consumeContext.Closed():
		workers.Wait()
		return errors.New("JetStream consumer stopped unexpectedly")
	}
}