392 lines
12 KiB
Go
392 lines
12 KiB
Go
// Package rocketmqx hides the Apache RocketMQ client behind the small surface
|
||
// this codebase needs: synchronous publish, delayed publish by millisecond
|
||
// timestamp, and clustering push consume.
|
||
package rocketmqx
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
rocketmq "github.com/apache/rocketmq-client-go/v2"
|
||
"github.com/apache/rocketmq-client-go/v2/consumer"
|
||
"github.com/apache/rocketmq-client-go/v2/primitive"
|
||
"github.com/apache/rocketmq-client-go/v2/producer"
|
||
)
|
||
|
||
const (
|
||
// Tencent Cloud RocketMQ accepts this property as an absolute Unix epoch
|
||
// millisecond delivery timestamp. Open-source RocketMQ delay levels remain
|
||
// available through the native client, but the rocket launch timer needs an
|
||
// exact backend-configured launch_at_ms.
|
||
PropertyStartDeliverTime = "__STARTDELIVERTIME"
|
||
)
|
||
|
||
// EndpointConfig is the shared RocketMQ connection identity.
|
||
type EndpointConfig struct {
|
||
NameServers []string
|
||
NameServerDomain string
|
||
AccessKey string
|
||
SecretKey string
|
||
SecurityToken string
|
||
Namespace string
|
||
VIPChannel bool
|
||
}
|
||
|
||
// ProducerConfig configures one producer group.
|
||
type ProducerConfig struct {
|
||
EndpointConfig
|
||
GroupName string
|
||
SendTimeout time.Duration
|
||
Retry int
|
||
}
|
||
|
||
// ConsumerConfig configures one clustering push consumer group.
|
||
type ConsumerConfig struct {
|
||
EndpointConfig
|
||
GroupName string
|
||
MaxReconsumeTimes int32
|
||
ConsumeRetryDelay time.Duration
|
||
ConsumePullBatch int32
|
||
ConsumePullTimeout time.Duration
|
||
ConsumeGoroutines int
|
||
// ConsumeFromFirst 仅给可重建读模型使用；普通业务 consumer 保持客户端默认的 last offset，避免新组误重放历史副作用。
|
||
ConsumeFromFirst bool
|
||
}
|
||
|
||
// Message is the stable publish shape used by services.
|
||
type Message struct {
|
||
Topic string
|
||
Tag string
|
||
Keys []string
|
||
Body []byte
|
||
Properties map[string]string
|
||
DeliveryAtMS int64
|
||
}
|
||
|
||
// ConsumedMessage is the stable consume shape used by service handlers.
|
||
type ConsumedMessage struct {
|
||
Topic string
|
||
Tag string
|
||
Keys string
|
||
MsgID string
|
||
Body []byte
|
||
Properties map[string]string
|
||
ReconsumeTimes int32
|
||
}
|
||
|
||
// Handler returns nil after the message has been durably processed. A non-nil
|
||
// error asks RocketMQ to retry the message according to the consumer group
|
||
// retry policy.
|
||
type Handler func(context.Context, ConsumedMessage) error
|
||
|
||
// Producer wraps a started or startable RocketMQ producer.
|
||
type Producer struct {
|
||
client rocketmq.Producer
|
||
}
|
||
|
||
// NewProducer creates a RocketMQ producer. Call Start before SendSync.
|
||
func NewProducer(cfg ProducerConfig) (*Producer, error) {
|
||
if err := validateEndpoint(cfg.EndpointConfig); err != nil {
|
||
return nil, err
|
||
}
|
||
group := strings.TrimSpace(cfg.GroupName)
|
||
if group == "" {
|
||
return nil, errors.New("rocketmq producer group_name is required")
|
||
}
|
||
options := []producer.Option{producer.WithGroupName(group)}
|
||
endpointOptions, err := producerEndpointOptions(cfg.EndpointConfig)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
options = append(options, endpointOptions...)
|
||
if cfg.SendTimeout > 0 {
|
||
options = append(options, producer.WithSendMsgTimeout(cfg.SendTimeout))
|
||
}
|
||
if cfg.Retry > 0 {
|
||
options = append(options, producer.WithRetry(cfg.Retry))
|
||
}
|
||
client, err := rocketmq.NewProducer(options...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &Producer{client: client}, nil
|
||
}
|
||
|
||
// Start registers the producer group and opens broker connections.
|
||
func (p *Producer) Start() error {
|
||
if p == nil || p.client == nil {
|
||
return errors.New("rocketmq producer is not configured")
|
||
}
|
||
return p.client.Start()
|
||
}
|
||
|
||
// Shutdown releases the underlying RocketMQ producer.
|
||
func (p *Producer) Shutdown() error {
|
||
if p == nil || p.client == nil {
|
||
return nil
|
||
}
|
||
return p.client.Shutdown()
|
||
}
|
||
|
||
// SendSync publishes one message and only returns nil when the broker accepts it.
|
||
func (p *Producer) SendSync(ctx context.Context, message Message) error {
|
||
if p == nil || p.client == nil {
|
||
return errors.New("rocketmq producer is not configured")
|
||
}
|
||
topic := strings.TrimSpace(message.Topic)
|
||
if topic == "" {
|
||
return errors.New("rocketmq topic is required")
|
||
}
|
||
msg := primitive.NewMessage(topic, message.Body)
|
||
if tag := strings.TrimSpace(message.Tag); tag != "" {
|
||
msg.WithTag(tag)
|
||
}
|
||
keys := make([]string, 0, len(message.Keys))
|
||
for _, key := range message.Keys {
|
||
if key = strings.TrimSpace(key); key != "" {
|
||
keys = append(keys, key)
|
||
}
|
||
}
|
||
if len(keys) > 0 {
|
||
msg.WithKeys(keys)
|
||
}
|
||
for key, value := range message.Properties {
|
||
msg.WithProperty(strings.TrimSpace(key), strings.TrimSpace(value))
|
||
}
|
||
if message.DeliveryAtMS > 0 {
|
||
msg.WithProperty(PropertyStartDeliverTime, fmt.Sprintf("%d", message.DeliveryAtMS))
|
||
}
|
||
result, err := p.client.SendSync(ctx, msg)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if result == nil || result.Status != primitive.SendOK {
|
||
if result == nil {
|
||
return errors.New("rocketmq send failed: empty result")
|
||
}
|
||
return fmt.Errorf("rocketmq send failed: status=%d msg_id=%s", result.Status, result.MsgID)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// Consumer wraps a clustering push consumer.
|
||
type Consumer struct {
|
||
mu sync.Mutex
|
||
client rocketmq.PushConsumer
|
||
started bool
|
||
}
|
||
|
||
// NewConsumer creates a RocketMQ push consumer. Subscribe before Start.
|
||
func NewConsumer(cfg ConsumerConfig) (*Consumer, error) {
|
||
if err := validateEndpoint(cfg.EndpointConfig); err != nil {
|
||
return nil, err
|
||
}
|
||
group := strings.TrimSpace(cfg.GroupName)
|
||
if group == "" {
|
||
return nil, errors.New("rocketmq consumer group_name is required")
|
||
}
|
||
options := []consumer.Option{consumer.WithGroupName(group)}
|
||
endpointOptions, err := consumerEndpointOptions(cfg.EndpointConfig)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
options = append(options, endpointOptions...)
|
||
if cfg.MaxReconsumeTimes > 0 {
|
||
options = append(options, consumer.WithMaxReconsumeTimes(cfg.MaxReconsumeTimes))
|
||
}
|
||
if cfg.ConsumeRetryDelay > 0 {
|
||
options = append(options, consumer.WithSuspendCurrentQueueTimeMillis(cfg.ConsumeRetryDelay))
|
||
}
|
||
if cfg.ConsumePullBatch > 0 {
|
||
options = append(options, consumer.WithPullBatchSize(cfg.ConsumePullBatch))
|
||
}
|
||
if cfg.ConsumePullTimeout > 0 {
|
||
options = append(options, consumer.WithConsumeTimeout(cfg.ConsumePullTimeout))
|
||
}
|
||
if cfg.ConsumeGoroutines > 0 {
|
||
options = append(options, consumer.WithConsumeGoroutineNums(cfg.ConsumeGoroutines))
|
||
}
|
||
if cfg.ConsumeFromFirst {
|
||
options = append(options, consumer.WithConsumeFromWhere(consumer.ConsumeFromFirstOffset))
|
||
}
|
||
client, err := rocketmq.NewPushConsumer(options...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &Consumer{client: client}, nil
|
||
}
|
||
|
||
// Subscribe binds a topic/tag selector to a handler. Empty tag subscribes all.
|
||
func (c *Consumer) Subscribe(topic string, tag string, handler Handler) error {
|
||
if c == nil || c.client == nil {
|
||
return errors.New("rocketmq consumer is not configured")
|
||
}
|
||
topic = strings.TrimSpace(topic)
|
||
if topic == "" {
|
||
return errors.New("rocketmq topic is required")
|
||
}
|
||
if handler == nil {
|
||
return errors.New("rocketmq handler is required")
|
||
}
|
||
selector := consumer.MessageSelector{}
|
||
if tag = strings.TrimSpace(tag); tag != "" {
|
||
selector = consumer.MessageSelector{Type: consumer.TAG, Expression: tag}
|
||
}
|
||
return c.client.Subscribe(topic, selector, func(ctx context.Context, messages ...*primitive.MessageExt) (consumer.ConsumeResult, error) {
|
||
for _, message := range messages {
|
||
if message == nil {
|
||
continue
|
||
}
|
||
err := handler(ctx, ConsumedMessage{
|
||
Topic: message.Topic,
|
||
Tag: message.GetTags(),
|
||
Keys: message.GetKeys(),
|
||
MsgID: message.MsgId,
|
||
Body: message.Body,
|
||
Properties: message.GetProperties(),
|
||
ReconsumeTimes: message.ReconsumeTimes,
|
||
})
|
||
if err != nil {
|
||
return consumer.ConsumeRetryLater, err
|
||
}
|
||
}
|
||
return consumer.ConsumeSuccess, nil
|
||
})
|
||
}
|
||
|
||
// Start begins pulling messages.
|
||
func (c *Consumer) Start() error {
|
||
if c == nil || c.client == nil {
|
||
return errors.New("rocketmq consumer is not configured")
|
||
}
|
||
c.mu.Lock()
|
||
defer c.mu.Unlock()
|
||
if c.started {
|
||
return nil
|
||
}
|
||
// RocketMQ 的 PushConsumer 只有 Start 成功后才具备完整的内部统计状态。
|
||
// 本地 namesrv 未启动时 Start 会返回错误，此时必须保持 started=false；
|
||
// 上层为了回滚已启动消费者会调用 Shutdown，未启动 consumer 直接跳过，避免依赖库在半初始化状态下 panic。
|
||
if err := c.client.Start(); err != nil {
|
||
return err
|
||
}
|
||
c.started = true
|
||
return nil
|
||
}
|
||
|
||
// Shutdown releases the underlying consumer.
|
||
func (c *Consumer) Shutdown() error {
|
||
if c == nil || c.client == nil {
|
||
return nil
|
||
}
|
||
c.mu.Lock()
|
||
defer c.mu.Unlock()
|
||
if !c.started {
|
||
return nil
|
||
}
|
||
defer func() {
|
||
// 无论底层关闭是否成功，本 wrapper 都不再重复调用 Shutdown；
|
||
// RocketMQ client 的二次关闭或半初始化关闭不具备幂等保证，重复尝试只会放大退出路径风险。
|
||
c.started = false
|
||
}()
|
||
return c.client.Shutdown()
|
||
}
|
||
|
||
func validateEndpoint(cfg EndpointConfig) error {
|
||
nameservers := normalizeStringSlice(cfg.NameServers)
|
||
if len(nameservers) == 0 && strings.TrimSpace(cfg.NameServerDomain) == "" {
|
||
return errors.New("rocketmq name_servers or name_server_domain is required")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func producerEndpointOptions(cfg EndpointConfig) ([]producer.Option, error) {
|
||
options := make([]producer.Option, 0, 5)
|
||
if nameservers := normalizeStringSlice(cfg.NameServers); len(nameservers) > 0 {
|
||
var err error
|
||
nameservers, err = resolveNameServers(nameservers)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
options = append(options, producer.WithNsResolver(primitive.NewPassthroughResolver(nameservers)))
|
||
} else if domain := strings.TrimSpace(cfg.NameServerDomain); domain != "" {
|
||
options = append(options, producer.WithNameServerDomain(domain))
|
||
}
|
||
if namespace := strings.TrimSpace(cfg.Namespace); namespace != "" {
|
||
options = append(options, producer.WithNamespace(namespace))
|
||
}
|
||
options = append(options, producer.WithVIPChannel(cfg.VIPChannel))
|
||
if strings.TrimSpace(cfg.AccessKey) != "" || strings.TrimSpace(cfg.SecretKey) != "" {
|
||
options = append(options, producer.WithCredentials(credentials(cfg)))
|
||
}
|
||
return options, nil
|
||
}
|
||
|
||
func consumerEndpointOptions(cfg EndpointConfig) ([]consumer.Option, error) {
|
||
options := make([]consumer.Option, 0, 5)
|
||
if nameservers := normalizeStringSlice(cfg.NameServers); len(nameservers) > 0 {
|
||
var err error
|
||
nameservers, err = resolveNameServers(nameservers)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
options = append(options, consumer.WithNsResolver(primitive.NewPassthroughResolver(nameservers)))
|
||
} else if domain := strings.TrimSpace(cfg.NameServerDomain); domain != "" {
|
||
options = append(options, consumer.WithNameServerDomain(domain))
|
||
}
|
||
if namespace := strings.TrimSpace(cfg.Namespace); namespace != "" {
|
||
options = append(options, consumer.WithNamespace(namespace))
|
||
}
|
||
options = append(options, consumer.WithVIPChannel(cfg.VIPChannel))
|
||
if strings.TrimSpace(cfg.AccessKey) != "" || strings.TrimSpace(cfg.SecretKey) != "" {
|
||
options = append(options, consumer.WithCredentials(credentials(cfg)))
|
||
}
|
||
return options, nil
|
||
}
|
||
|
||
func resolveNameServers(nameservers []string) ([]string, error) {
|
||
resolved := make([]string, 0, len(nameservers))
|
||
for _, nameserver := range nameservers {
|
||
host, port, err := net.SplitHostPort(nameserver)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("invalid rocketmq name_server %q: %w", nameserver, err)
|
||
}
|
||
if net.ParseIP(host) != nil {
|
||
resolved = append(resolved, net.JoinHostPort(host, port))
|
||
continue
|
||
}
|
||
ips, err := net.LookupHost(host)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("resolve rocketmq name_server %q: %w", nameserver, err)
|
||
}
|
||
if len(ips) == 0 {
|
||
return nil, fmt.Errorf("resolve rocketmq name_server %q: empty result", nameserver)
|
||
}
|
||
resolved = append(resolved, net.JoinHostPort(ips[0], port))
|
||
}
|
||
return resolved, nil
|
||
}
|
||
|
||
func credentials(cfg EndpointConfig) primitive.Credentials {
|
||
return primitive.Credentials{
|
||
AccessKey: strings.TrimSpace(cfg.AccessKey),
|
||
SecretKey: strings.TrimSpace(cfg.SecretKey),
|
||
SecurityToken: strings.TrimSpace(cfg.SecurityToken),
|
||
}
|
||
}
|
||
|
||
func normalizeStringSlice(input []string) []string {
|
||
out := make([]string, 0, len(input))
|
||
for _, value := range input {
|
||
if value = strings.TrimSpace(value); value != "" {
|
||
out = append(out, value)
|
||
}
|
||
}
|
||
return out
|
||
}
|