lislgosms/gateway/internal/submitworker/worker.go

package submitworker

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"cmpp-platform/gateway/internal/metrics"
	"cmpp-platform/gateway/internal/queue"
	"cmpp-platform/gateway/internal/ratelimit"
	"cmpp-platform/gateway/internal/upstream"

	"github.com/redis/go-redis/v9"
)

const (
	defaultStream      = "gateway.submit.commands"
	defaultGroup       = "cmpp-gateway"
	defaultConsumer    = "gateway-1"
	defaultMinIdle     = 30 * time.Second
	defaultMaxFails    = 3
	defaultConcurrency = 64
)

type Worker struct {
	Redis            *redis.Client
	Upstream         *upstream.Manager
	Limiter          ratelimit.Limiter
	Submit           func(context.Context, queue.SubmitCommand) (queue.SubmitResult, error)
	ReportDeadLetter func(context.Context, DeadLetterEvent) error
	Stream           string
	Group            string
	Consumer         string
	Block            time.Duration
	Count            int64
	Concurrency      int
	MinIdle          time.Duration
	MaxFailures      int
	APIBaseURL       string
	HTTPClient       *http.Client
	Logger           *log.Logger
	inFlight         atomic.Int64
}

type DeadLetterEvent struct {
	StreamMessageID string                 `json:"streamMessageId"`
	TraceID         string                 `json:"traceId,omitempty"`
	MessageID       string                 `json:"messageId,omitempty"`
	ChannelID       string                 `json:"channelId,omitempty"`
	TenantID        string                 `json:"tenantId,omitempty"`
	ApplicationID   string                 `json:"applicationId,omitempty"`
	SubmitID        string                 `json:"submitId,omitempty"`
	FailureCode     string                 `json:"failureCode"`
	FailureMessage  string                 `json:"failureMessage"`
	Attempts        int                    `json:"attempts"`
	MaxAttempts     int                    `json:"maxAttempts"`
	CommandPayload  map[string]interface{} `json:"commandPayload,omitempty"`
	RawPayload      string                 `json:"rawPayload,omitempty"`
	DeadLetteredAt  time.Time              `json:"deadLetteredAt"`
}

func New(redisURL string, manager *upstream.Manager) (*Worker, error) {
	client, err := redisClient(redisURL)
	if err != nil {
		return nil, err
	}
	return &Worker{Redis: client, Upstream: manager, Limiter: ratelimit.NewWithClient(client)}, nil
}

func (w *Worker) Run(ctx context.Context) error {
	if w.Redis == nil {
		return fmt.Errorf("redis client is required")
	}
	if w.Upstream == nil {
		return fmt.Errorf("upstream manager is required")
	}
	pool := newMessageWorkPool(ctx, w, w.concurrency())
	defer pool.wait()
	for {
		if err := w.ensureGroup(ctx); err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			w.logf("gateway submit worker ensure group failed: %v", err)
			sleep(ctx, 3*time.Second)
			continue
		}
		if err := w.recoverPending(ctx, pool); err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			w.logf("gateway submit worker pending recovery failed: %v", err)
			sleep(ctx, time.Second)
			continue
		}
		if err := w.consumeOnce(ctx, pool); err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			w.logf("gateway submit worker consume failed: %v", err)
			sleep(ctx, time.Second)
		}
	}
}

func (w *Worker) ensureGroup(ctx context.Context) error {
	err := w.Redis.XGroupCreateMkStream(ctx, w.stream(), w.group(), "0").Err()
	if err == nil || strings.Contains(err.Error(), "BUSYGROUP") {
		return nil
	}
	return err
}

func (w *Worker) consumeOnce(ctx context.Context, pool *messageWorkPool) error {
	available, err := pool.waitForCapacity(ctx)
	if err != nil {
		return err
	}
	streams, err := w.Redis.XReadGroup(ctx, &redis.XReadGroupArgs{
		Group:    w.group(),
		Consumer: w.consumer(),
		Streams:  []string{w.stream(), ">"},
		Count:    min(w.count(), int64(available)),
		Block:    w.block(),
	}).Result()
	if errors.Is(err, redis.Nil) {
		return nil
	}
	if err != nil {
		return err
	}
	for _, stream := range streams {
		for _, message := range stream.Messages {
			if !pool.dispatch(message) {
				return fmt.Errorf("gateway submit worker capacity accounting mismatch")
			}
		}
	}
	return nil
}

func (w *Worker) recoverPending(ctx context.Context, pool *messageWorkPool) error {
	start := "0-0"
	for {
		available := pool.available()
		if available == 0 {
			return nil
		}
		messages, next, err := w.Redis.XAutoClaim(ctx, &redis.XAutoClaimArgs{
			Stream:   w.stream(),
			Group:    w.group(),
			Consumer: w.consumer(),
			MinIdle:  w.minIdle(),
			Start:    start,
			Count:    min(w.count(), int64(available)),
		}).Result()
		if errors.Is(err, redis.Nil) {
			return nil
		}
		if err != nil {
			return err
		}
		if len(messages) == 0 {
			return nil
		}
		w.logf("gateway submit worker reclaimed %d pending message(s)", len(messages))
		for _, message := range messages {
			// An in-flight command can legitimately exceed MinIdle while waiting on a supplier.
			// Rechecking both the local active set and Redis PEL closes the race where the
			// original attempt ACKs between XAUTOCLAIM returning and local dispatch.
			if err := pool.dispatchRecovered(ctx, message); err != nil {
				return err
			}
		}
		start = next
		if next == "0-0" {
			return nil
		}
	}
}

type messageWorkPool struct {
	ctx       context.Context
	worker    *Worker
	slots     chan struct{}
	completed chan struct{}
	group     sync.WaitGroup
	mu        sync.Mutex
	active    map[string]struct{}
}

func newMessageWorkPool(ctx context.Context, worker *Worker, concurrency int) *messageWorkPool {
	return &messageWorkPool{
		ctx: ctx, worker: worker, slots: make(chan struct{}, concurrency),
		completed: make(chan struct{}, concurrency), active: make(map[string]struct{}),
	}
}

func (p *messageWorkPool) available() int {
	return cap(p.slots) - len(p.slots)
}

func (p *messageWorkPool) waitForCapacity(ctx context.Context) (int, error) {
	for p.available() == 0 {
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		case <-p.completed:
		}
	}
	return p.available(), nil
}

func (p *messageWorkPool) dispatch(message redis.XMessage) bool {
	p.mu.Lock()
	if _, exists := p.active[message.ID]; exists {
		p.mu.Unlock()
		return true
	}
	select {
	case p.slots <- struct{}{}:
		p.active[message.ID] = struct{}{}
		p.worker.inFlight.Add(1)
		p.group.Add(1)
		p.mu.Unlock()
	case <-p.ctx.Done():
		p.mu.Unlock()
		return false
	default:
		p.mu.Unlock()
		return false
	}
	go func() {
		defer func() {
			p.mu.Lock()
			delete(p.active, message.ID)
			p.mu.Unlock()
			<-p.slots
			p.worker.inFlight.Add(-1)
			select {
			case p.completed <- struct{}{}:
			default:
			}
			p.group.Done()
		}()
		if err := p.worker.processMessage(p.ctx, message); err != nil {
			p.worker.logf("gateway submit worker message %s failed: %v", message.ID, err)
		}
	}()
	return true
}

func (p *messageWorkPool) dispatchRecovered(ctx context.Context, message redis.XMessage) error {
	p.mu.Lock()
	_, active := p.active[message.ID]
	p.mu.Unlock()
	if active {
		return nil
	}
	pending, err := p.worker.Redis.XPendingExt(ctx, &redis.XPendingExtArgs{
		Stream: p.worker.stream(), Group: p.worker.group(), Start: message.ID, End: message.ID, Count: 1,
	}).Result()
	if err != nil && !errors.Is(err, redis.Nil) {
		return err
	}
	if len(pending) == 0 {
		return nil
	}
	if !p.dispatch(message) {
		return fmt.Errorf("gateway submit worker recovery capacity accounting mismatch")
	}
	return nil
}

func (p *messageWorkPool) wait() {
	p.group.Wait()
}

func (w *Worker) processMessage(ctx context.Context, message redis.XMessage) error {
	command, err := CommandFromStreamValues(message.Values)
	if err != nil {
		return w.deadLetterMalformedMessage(ctx, message, err)
	}
	if !command.CreatedAt.IsZero() {
		metrics.ObserveSubmitStage("stream_wait", true, time.Since(command.CreatedAt))
	}
	if err := w.handleCommand(ctx, command); err != nil {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		attempts, attemptsErr := w.incrementFailureAttempt(ctx, message.ID)
		if attemptsErr != nil {
			w.logf("gateway submit worker increment failure %s failed: %v", message.ID, attemptsErr)
		}
		if attempts >= w.maxFailures() {
			if reportErr := w.deadLetterCommand(ctx, message, command, attempts, err); reportErr != nil {
				return reportErr
			}
			return w.ackAndClearFailure(ctx, message.ID)
		}
		return err
	}
	return w.ackAndClearFailure(ctx, message.ID)
}

func (w *Worker) handleCommand(ctx context.Context, command queue.SubmitCommand) error {
	startedAt := time.Now()
	limitStartedAt := time.Now()
	if w.Limiter != nil {
		if _, err := w.Limiter.Wait(ctx, command.ChannelID, command.Route.RateLimitPerSecond); err != nil {
			metrics.ObserveSubmitStage("rate_limit_wait", false, time.Since(limitStartedAt))
			return err
		}
	}
	metrics.ObserveSubmitStage("rate_limit_wait", true, time.Since(limitStartedAt))
	submit := w.Submit
	if submit == nil {
		if w.Upstream == nil {
			return fmt.Errorf("upstream manager is required")
		}
		submit = w.Upstream.Submit
	}
	result, err := submit(ctx, command)
	accepted := err == nil && (result.SubmitStatus == "" || result.SubmitStatus == "accepted")
	metrics.ObserveSubmit(accepted, time.Since(startedAt))
	if err != nil && (result.SubmitStatus == "" || result.SubmitStatus == "accepted") {
		return err
	}
	return nil
}

func CommandFromStreamValues(values map[string]interface{}) (queue.SubmitCommand, error) {
	raw, ok := values["data"]
	if !ok {
		return queue.SubmitCommand{}, fmt.Errorf("stream data field is required")
	}
	var data string
	switch value := raw.(type) {
	case string:
		data = value
	case []byte:
		data = string(value)
	default:
		data = fmt.Sprint(value)
	}
	var command queue.SubmitCommand
	if err := json.Unmarshal([]byte(data), &command); err != nil {
		return queue.SubmitCommand{}, err
	}
	if command.MessageType != queue.MessageTypeSubmitCommand {
		return queue.SubmitCommand{}, fmt.Errorf("unsupported messageType %q", command.MessageType)
	}
	return command, nil
}

func (w *Worker) deadLetterMalformedMessage(ctx context.Context, message redis.XMessage, cause error) error {
	event := DeadLetterEvent{
		StreamMessageID: message.ID,
		FailureCode:     "INVALID_COMMAND_PAYLOAD",
		FailureMessage:  cause.Error(),
		Attempts:        1,
		MaxAttempts:     1,
		RawPayload:      extractRawPayload(message.Values),
		DeadLetteredAt:  time.Now().UTC(),
	}
	if err := w.reportDeadLetter(ctx, event); err != nil {
		return err
	}
	return w.ackAndClearFailure(ctx, message.ID)
}

func (w *Worker) deadLetterCommand(ctx context.Context, message redis.XMessage, command queue.SubmitCommand, attempts int, cause error) error {
	payload, payloadErr := commandPayload(command)
	if payloadErr != nil {
		w.logf("gateway submit worker marshal dead-letter command %s failed: %v", message.ID, payloadErr)
	}
	return w.reportDeadLetter(ctx, DeadLetterEvent{
		StreamMessageID: message.ID,
		TraceID:         command.TraceID,
		MessageID:       command.MessageID,
		ChannelID:       command.ChannelID,
		TenantID:        command.TenantID,
		ApplicationID:   command.ApplicationID,
		SubmitID:        command.SubmitID,
		FailureCode:     "SUBMIT_PROCESSING_FAILED",
		FailureMessage:  cause.Error(),
		Attempts:        attempts,
		MaxAttempts:     w.maxFailures(),
		CommandPayload:  payload,
		RawPayload:      extractRawPayload(message.Values),
		DeadLetteredAt:  time.Now().UTC(),
	})
}

func (w *Worker) reportDeadLetter(ctx context.Context, event DeadLetterEvent) error {
	if w.ReportDeadLetter != nil {
		return w.ReportDeadLetter(ctx, event)
	}
	if w.APIBaseURL == "" {
		return fmt.Errorf("gateway submit dead-letter reporter is not configured")
	}
	client := w.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: 10 * time.Second}
	}
	body, err := json.Marshal(event)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		strings.TrimRight(w.APIBaseURL, "/")+"/gateway/events/dead-letter",
		bytes.NewReader(body),
	)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("dead-letter endpoint returned %d", resp.StatusCode)
	}
	return nil
}

func redisClient(redisURL string) (*redis.Client, error) {
	if redisURL == "" {
		redisURL = "redis://127.0.0.1:6379"
	}
	options, err := redis.ParseURL(redisURL)
	if err != nil {
		return nil, err
	}
	return redis.NewClient(options), nil
}

func (w *Worker) incrementFailureAttempt(ctx context.Context, messageID string) (int, error) {
	value, err := w.Redis.HIncrBy(ctx, w.failureAttemptsKey(), messageID, 1).Result()
	return int(value), err
}

func (w *Worker) ackAndClearFailure(ctx context.Context, messageID string) error {
	if err := w.Redis.XAck(ctx, w.stream(), w.group(), messageID).Err(); err != nil {
		return err
	}
	if err := w.Redis.HDel(ctx, w.failureAttemptsKey(), messageID).Err(); err != nil {
		w.logf("gateway submit worker clear failure %s failed: %v", messageID, err)
	}
	return nil
}

func (w *Worker) stream() string {
	if w.Stream != "" {
		return w.Stream
	}
	return defaultStream
}

func (w *Worker) group() string {
	if w.Group != "" {
		return w.Group
	}
	return defaultGroup
}

func (w *Worker) consumer() string {
	if w.Consumer != "" {
		return w.Consumer
	}
	return defaultConsumer
}

func (w *Worker) block() time.Duration {
	if w.Block > 0 {
		return w.Block
	}
	return 5 * time.Second
}

func (w *Worker) count() int64 {
	if w.Count > 0 {
		return w.Count
	}
	return 10
}

func (w *Worker) concurrency() int {
	if w.Concurrency > 0 {
		return min(w.Concurrency, 1024)
	}
	return defaultConcurrency
}

func (w *Worker) ConfiguredConcurrency() int {
	return w.concurrency()
}

func (w *Worker) InFlight() int64 {
	return w.inFlight.Load()
}

func (w *Worker) minIdle() time.Duration {
	if w.MinIdle > 0 {
		return w.MinIdle
	}
	return defaultMinIdle
}

func (w *Worker) maxFailures() int {
	if w.MaxFailures > 0 {
		return w.MaxFailures
	}
	return defaultMaxFails
}

func (w *Worker) failureAttemptsKey() string {
	return w.stream() + ":failure-attempts"
}

func (w *Worker) logf(format string, args ...interface{}) {
	if w.Logger != nil {
		w.Logger.Printf(format, args...)
		return
	}
	log.Printf(format, args...)
}

func sleep(ctx context.Context, duration time.Duration) {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
	case <-timer.C:
	}
}

func extractRawPayload(values map[string]interface{}) string {
	raw, ok := values["data"]
	if !ok {
		return ""
	}
	switch value := raw.(type) {
	case string:
		return value
	case []byte:
		return string(value)
	default:
		return fmt.Sprint(value)
	}
}

func commandPayload(command queue.SubmitCommand) (map[string]interface{}, error) {
	data, err := json.Marshal(command)
	if err != nil {
		return nil, err
	}
	var payload map[string]interface{}
	if err := json.Unmarshal(data, &payload); err != nil {
		return nil, err
	}
	return payload, nil
}