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 }