yovision/Sense/server/app/sense/outbox/relay.go

package outbox

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

var (
	ErrDuplicateIdempotency = errors.New("可靠投递幂等键已存在")
	ErrLeaseLost            = errors.New("可靠投递租约已失效")
)

func Enqueue(tx *gorm.DB, input EnqueueInput, now time.Time) (Message, error) {
	input.InternalType, input.BusinessRef, input.IdempotencyKey = strings.TrimSpace(input.InternalType), strings.TrimSpace(input.BusinessRef), strings.TrimSpace(input.IdempotencyKey)
	if input.InternalType == "" || len(input.InternalType) > 64 || input.BusinessRef == "" || len(input.BusinessRef) > 128 || input.IdempotencyKey == "" || len(input.IdempotencyKey) > 191 || !json.Valid(input.PayloadJSON) {
		return Message{}, ErrInvalidInput
	}
	if input.MaxAttempts == 0 {
		input.MaxAttempts = 12
	}
	if input.MaxAttempts < 1 || input.MaxAttempts > 100 {
		return Message{}, ErrInvalidInput
	}
	now = now.UTC()
	message := Message{ID: uuid.NewString(), InternalType: input.InternalType, BusinessRef: input.BusinessRef, IdempotencyKey: input.IdempotencyKey, PayloadJSON: string(input.PayloadJSON), State: StatePending, MaxAttempts: input.MaxAttempts, AvailableAt: now, Version: 1, CreatedAt: now, UpdatedAt: now}
	var existing int64
	if err := tx.Model(&Message{}).Where("idempotency_key = ?", input.IdempotencyKey).Count(&existing).Error; err != nil {
		return Message{}, fmt.Errorf("check outbox idempotency: %w", err)
	}
	if existing > 0 {
		return Message{}, ErrDuplicateIdempotency
	}
	if err := tx.Create(&message).Error; err != nil {
		return Message{}, fmt.Errorf("enqueue outbox message: %w", err)
	}
	return message, nil
}

type Relay struct {
	DB            *gorm.DB
	Now           func() time.Time
	LeaseDuration time.Duration
	Backoff       func(int) time.Duration
}

func NewRelay(db *gorm.DB) *Relay {
	return &Relay{DB: db, Now: time.Now, LeaseDuration: 30 * time.Second, Backoff: defaultBackoff}
}

func (r *Relay) Claim(worker string, limit int) ([]Message, error) {
	worker = strings.TrimSpace(worker)
	if worker == "" || len(worker) > 128 || limit < 1 || limit > 100 {
		return nil, ErrInvalidInput
	}
	now, leaseUntil := r.now(), r.now().Add(r.leaseDuration())
	claimed := make([]Message, 0, limit)
	err := r.DB.Transaction(func(tx *gorm.DB) error {
		var candidates []Message
		query := tx.Where("((state IN ?) AND available_at <= ?) OR (state = ? AND lease_until < ?)", []string{StatePending, StateRetry}, now, StateProcessing, now).Order("available_at ASC, created_at ASC").Limit(limit)
		if tx.Dialector.Name() == "postgres" {
			query = query.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"})
		}
		if err := query.Find(&candidates).Error; err != nil {
			return err
		}
		for _, item := range candidates {
			result := tx.Model(&Message{}).Where("id = ? AND version = ?", item.ID, item.Version).Updates(map[string]interface{}{"state": StateProcessing, "lease_owner": worker, "lease_until": leaseUntil, "version": gorm.Expr("version + 1"), "updated_at": now})
			if result.Error != nil {
				return result.Error
			}
			if result.RowsAffected == 1 {
				item.State, item.LeaseOwner, item.LeaseUntil, item.Version = StateProcessing, worker, &leaseUntil, item.Version+1
				claimed = append(claimed, item)
			}
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("claim outbox messages: %w", err)
	}
	return claimed, nil
}

func (r *Relay) MarkSuccess(id, worker string) error {
	now := r.now()
	return r.DB.Transaction(func(tx *gorm.DB) error {
		var message Message
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&message, "id = ?", id).Error; err != nil {
			return err
		}
		var delivered int64
		if err := tx.Model(&DeliveryRecord{}).Where("idempotency_key = ?", message.IdempotencyKey).Count(&delivered).Error; err != nil {
			return err
		}
		if delivered > 0 {
			return nil
		}
		if message.State != StateProcessing || message.LeaseOwner != worker || message.LeaseUntil == nil || !message.LeaseUntil.After(now) {
			return ErrLeaseLost
		}
		if err := tx.Create(&DeliveryRecord{IdempotencyKey: message.IdempotencyKey, MessageID: message.ID, DeliveredAt: now, CreatedAt: now}).Error; err != nil {
			return err
		}
		number := message.AttemptCount + 1
		if err := tx.Create(&Attempt{MessageID: message.ID, Number: number, Outcome: StateDelivered, Detail: "投递成功", Worker: worker, CreatedAt: now}).Error; err != nil {
			return err
		}
		return tx.Model(&message).Updates(map[string]interface{}{"state": StateDelivered, "attempt_count": number, "delivered_at": now, "lease_owner": "", "lease_until": nil, "last_error": "", "version": gorm.Expr("version + 1"), "updated_at": now}).Error
	})
}

func (r *Relay) MarkFailure(id, worker, detail string) error {
	now := r.now()
	detail = strings.TrimSpace(detail)
	if len(detail) > 512 {
		detail = detail[:512]
	}
	if detail == "" {
		detail = "投递失败"
	}
	return r.DB.Transaction(func(tx *gorm.DB) error {
		var message Message
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&message, "id = ?", id).Error; err != nil {
			return err
		}
		if message.State != StateProcessing || message.LeaseOwner != worker || message.LeaseUntil == nil || !message.LeaseUntil.After(now) {
			return ErrLeaseLost
		}
		nextAttempt := message.AttemptCount + 1
		state := StateRetry
		available := now.Add(r.backoff(nextAttempt))
		if nextAttempt >= message.MaxAttempts {
			state = StateDead
			available = now
		}
		if err := tx.Model(&message).Updates(map[string]interface{}{"state": state, "attempt_count": nextAttempt, "available_at": available, "lease_owner": "", "lease_until": nil, "last_error": detail, "version": gorm.Expr("version + 1"), "updated_at": now}).Error; err != nil {
			return err
		}
		return tx.Create(&Attempt{MessageID: message.ID, Number: nextAttempt, Outcome: state, Detail: detail, Worker: worker, CreatedAt: now}).Error
	})
}

func (r *Relay) now() time.Time {
	if r.Now != nil {
		return r.Now().UTC()
	}
	return time.Now().UTC()
}
func (r *Relay) leaseDuration() time.Duration {
	if r.LeaseDuration <= 0 {
		return 30 * time.Second
	}
	return r.LeaseDuration
}
func (r *Relay) backoff(attempt int) time.Duration {
	if r.Backoff != nil {
		return r.Backoff(attempt)
	}
	return defaultBackoff(attempt)
}
func defaultBackoff(attempt int) time.Duration {
	if attempt < 1 {
		attempt = 1
	}
	delay := time.Second * time.Duration(1<<min(attempt-1, 8))
	if delay > 5*time.Minute {
		return 5 * time.Minute
	}
	return delay
}