iotlcm/internal/consumer/rabbitmq/batch.go

package rabbitmq

import (
	"context"
	"encoding/json"
	"errors"
	"strconv"
	"time"

	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	amqp "github.com/rabbitmq/amqp091-go"

	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/batch"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/consts"
	mq "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/library/rabbitmq"
)

const batchRetryCountHeader = "x-iotlcm-retry-count"

func (c *Consumers) initBatchExecute(ctx context.Context) error {
	if !mq.Enabled(ctx) {
		g.Log().Info(ctx, "batch execute consumer skipped, RabbitMQ disabled")
		return nil
	}
	if c.batch == nil {
		return gerror.New("batch execute handler is not configured")
	}
	retryDelay, maxRetries, err := getBatchExecuteRetryConfig(ctx)
	if err != nil {
		return err
	}
	c.batchExecuteRetryDelay = retryDelay
	c.batchExecuteMaxRetries = maxRetries
	if err = mq.DeclareQueue(ctx, consts.QueueBatchExecuteDeadLetter); err != nil {
		return gerror.Wrap(err, "declare batch execute dead-letter queue failed")
	}
	if err = mq.ConsumeJSON[batch.Message](
		ctx,
		consts.QueueBatchExecute,
		"iotlcm-batch-execute",
		c.handleBatchExecute,
	); err != nil {
		return err
	}
	g.Log().Infof(ctx, "batch execute consumer started, queue:%s, retryDelay:%s, maxRetries:%d", consts.QueueBatchExecute, retryDelay, maxRetries)
	return nil
}

func (c *Consumers) handleBatchExecute(ctx context.Context, message batch.Message) error {
	if c.batch == nil {
		return gerror.New("batch execute handler is not configured")
	}
	g.Log().Infof(ctx, "batch execute message received, batchNo:%s, batchType:%d", message.BatchNo, message.BatchType)
	if err := c.batch.Handle(ctx, message); err != nil {
		if _, delivered := mq.DeliveryMetadataFromContext(ctx); !delivered {
			return err
		}
		if batch.IsRetryable(err) {
			message.Retryable = true
		}
		var failureMarkErr *batch.FailureMarkError
		if errors.As(err, &failureMarkErr) {
			message.FailureReason = failureMarkErr.Reason
		}
		return c.retryOrDeadLetterBatchMessage(ctx, message, err)
	}
	return nil
}

func (c *Consumers) retryOrDeadLetterBatchMessage(ctx context.Context, message batch.Message, handlerErr error) error {
	metadata, _ := mq.DeliveryMetadataFromContext(ctx)
	retryCount := batchRetryCount(metadata.Headers)
	headers := amqp.Table{batchRetryCountHeader: retryCount + 1}
	if shouldDeadLetterBatch(retryCount, c.batchExecuteMaxRetries) {
		if message.Retryable {
			finalizationMessage := message
			finalizationMessage.Retryable = false
			finalizationMessage.FailureReason = handlerErr.Error()
			if err := c.batch.Handle(ctx, finalizationMessage); err != nil {
				return gerror.Wrap(err, "finalize retryable batch failure before dead-lettering")
			}
			message.FailureReason = handlerErr.Error()
		}
		deadLetterBody, marshalErr := json.Marshal(batchDeadLetterMessage{
			Message:    message,
			RetryCount: retryCount,
			Error:      handlerErr.Error(),
			FailedAt:   time.Now().Format(time.RFC3339),
		})
		if marshalErr != nil {
			return gerror.Wrap(marshalErr, "marshal batch execute dead-letter message failed")
		}
		if err := mq.PublishBytesWithHeaders(ctx, consts.QueueBatchExecuteDeadLetter, deadLetterBody, nil); err != nil {
			return gerror.Wrap(err, "publish batch execute dead-letter message failed")
		}
		g.Log().Errorf(ctx, "batch execute message sent to dead-letter queue, batchNo:%s, batchType:%d, retryCount:%d, err:%+v", message.BatchNo, message.BatchType, retryCount, handlerErr)
		return nil
	}
	body, marshalErr := json.Marshal(message)
	if marshalErr != nil {
		return gerror.Wrap(marshalErr, "marshal batch execute retry message failed")
	}
	if err := mq.PublishDelayedBytesWithHeaders(
		ctx,
		consts.QueueBatchExecuteRetry,
		consts.QueueBatchExecute,
		c.batchExecuteRetryDelay,
		body,
		headers,
	); err != nil {
		return gerror.Wrap(err, "republish batch execute retry message failed")
	}
	g.Log().Warningf(ctx, "batch execute message scheduled for retry, batchNo:%s, batchType:%d, retryCount:%d, delay:%s, err:%+v", message.BatchNo, message.BatchType, retryCount+1, c.batchExecuteRetryDelay, handlerErr)
	return nil
}

type batchDeadLetterMessage struct {
	Message    batch.Message `json:"message"`
	RetryCount int           `json:"retryCount"`
	Error      string        `json:"error"`
	FailedAt   string        `json:"failedAt"`
}

func shouldDeadLetterBatch(retryCount, maxRetries int) bool {
	return retryCount >= maxRetries
}

func batchRetryCount(headers map[string]any) int {
	if headers == nil {
		return 0
	}
	value, ok := headers[batchRetryCountHeader]
	if !ok {
		return 0
	}
	switch count := value.(type) {
	case int:
		return max(count, 0)
	case int8:
		return max(int(count), 0)
	case int16:
		return max(int(count), 0)
	case int32:
		return max(int(count), 0)
	case int64:
		return max(int(count), 0)
	case uint8:
		return int(count)
	case uint16:
		return int(count)
	case uint32:
		return int(count)
	case uint64:
		if count > uint64(^uint(0)>>1) {
			return int(^uint(0) >> 1)
		}
		return int(count)
	case string:
		parsed, err := strconv.Atoi(count)
		if err == nil {
			return max(parsed, 0)
		}
	}
	return 0
}

func getBatchExecuteRetryConfig(ctx context.Context) (time.Duration, int, error) {
	delay := g.Cfg().MustGet(ctx, "rabbitmq.batchExecute.retryDelay", "30s").Duration()
	if delay <= 0 {
		return 0, 0, gerror.New("batch execute retry delay must be positive")
	}
	maxRetries := g.Cfg().MustGet(ctx, "rabbitmq.batchExecute.maxRetries", 3).Int()
	if maxRetries < 0 {
		return 0, 0, gerror.New("batch execute max retries cannot be negative")
	}
	return delay, maxRetries, nil
}