362 lines
11 KiB
Go
362 lines
11 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"
|
|
"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
|
|
}
|
|
|
|
// 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 {
|
|
client rocketmq.PushConsumer
|
|
}
|
|
|
|
// 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))
|
|
}
|
|
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")
|
|
}
|
|
return c.client.Start()
|
|
}
|
|
|
|
// Shutdown releases the underlying consumer.
|
|
func (c *Consumer) Shutdown() error {
|
|
if c == nil || c.client == nil {
|
|
return nil
|
|
}
|
|
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
|
|
}
|