97 lines
2.5 KiB
Go
97 lines
2.5 KiB
Go
// Package assignmentqueue implements the durable assignment handoff with JetStream.
|
|
package assignmentqueue
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"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 Handler interface {
|
|
Handle(context.Context, taskassignment.Assignment) (bool, error)
|
|
}
|
|
|
|
// Message is the subset of jetstream.Msg needed by one reconciliation.
|
|
type Message interface {
|
|
Data() []byte
|
|
DoubleAck(context.Context) error
|
|
NakWithDelay(time.Duration) error
|
|
InProgress() error
|
|
TermWithReason(string) error
|
|
}
|
|
|
|
// Processor maps one delivery to one idempotent worker reconciliation.
|
|
type Processor struct {
|
|
TrustDomain string
|
|
Handler Handler
|
|
RetryDelay time.Duration
|
|
}
|
|
|
|
func (p Processor) Process(ctx context.Context, message Message) error {
|
|
if p.Handler == nil {
|
|
return errors.New("assignment handler is required")
|
|
}
|
|
assignment, err := taskassignment.Unmarshal(message.Data(), p.TrustDomain)
|
|
if err != nil {
|
|
return errors.Join(err, message.TermWithReason("invalid assignment"))
|
|
}
|
|
done, err := p.Handler.Handle(ctx, assignment)
|
|
if err != nil {
|
|
delay := p.RetryDelay
|
|
if delay <= 0 {
|
|
delay = 15 * time.Second
|
|
}
|
|
return errors.Join(err, message.NakWithDelay(delay))
|
|
}
|
|
if done {
|
|
if err := message.DoubleAck(ctx); err != nil {
|
|
return fmt.Errorf("ack assignment %s: %w", assignment.ID, err)
|
|
}
|
|
return nil
|
|
}
|
|
if err := message.InProgress(); err != nil {
|
|
return fmt.Errorf("extend assignment %s acknowledgement: %w", assignment.ID, err)
|
|
}
|
|
return nil
|
|
}
|