migrate from github
This commit is contained in:
@@ -0,0 +1,86 @@
|
||||
package mq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/pkg/errs"
|
||||
)
|
||||
|
||||
type Message struct {
|
||||
ID string
|
||||
Topic string
|
||||
Key string
|
||||
Body []byte
|
||||
Headers map[string]string
|
||||
Timestamp time.Time
|
||||
}
|
||||
|
||||
type PublishMessage struct {
|
||||
Topic string
|
||||
Key string
|
||||
Body []byte
|
||||
Headers map[string]string
|
||||
}
|
||||
|
||||
type ConsumeMode string
|
||||
|
||||
const (
|
||||
ConsumeModeFanout ConsumeMode = "fanout"
|
||||
ConsumeModeGroup ConsumeMode = "group"
|
||||
)
|
||||
|
||||
type ErrorAction string
|
||||
|
||||
const (
|
||||
ErrorActionDrop ErrorAction = "drop"
|
||||
ErrorActionRequeue ErrorAction = "requeue"
|
||||
)
|
||||
|
||||
type ConsumeOptions struct {
|
||||
Topic string
|
||||
Mode ConsumeMode
|
||||
Group string
|
||||
Consumer string
|
||||
AutoAck bool
|
||||
OnError ErrorAction
|
||||
}
|
||||
|
||||
func (o ConsumeOptions) normalize() ConsumeOptions {
|
||||
n := o
|
||||
n.Topic = strings.TrimSpace(n.Topic)
|
||||
n.Group = strings.TrimSpace(n.Group)
|
||||
n.Consumer = strings.TrimSpace(n.Consumer)
|
||||
if n.Mode == "" {
|
||||
n.Mode = ConsumeModeFanout
|
||||
}
|
||||
if n.OnError == "" {
|
||||
n.OnError = ErrorActionRequeue
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (o ConsumeOptions) validate() error {
|
||||
switch o.Mode {
|
||||
case ConsumeModeFanout, ConsumeModeGroup:
|
||||
default:
|
||||
return errs.Detailf(errs.ErrInvalidConsumeMode, "%s", o.Mode)
|
||||
}
|
||||
|
||||
switch o.OnError {
|
||||
case ErrorActionDrop, ErrorActionRequeue:
|
||||
default:
|
||||
return errs.Detailf(errs.ErrInvalidErrorAction, "%s", o.OnError)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type MessageHandler func(ctx context.Context, msg Message) error
|
||||
|
||||
type MessageQueue interface {
|
||||
Publish(ctx context.Context, msg PublishMessage) error
|
||||
Consume(ctx context.Context, opts ConsumeOptions, handler MessageHandler) error
|
||||
Close() error
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package mq
|
||||
|
||||
import "strings"
|
||||
|
||||
type Backend string
|
||||
|
||||
const (
|
||||
BackendRedis Backend = "redis"
|
||||
BackendRabbitMQ Backend = "rabbitmq"
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
RabbitMQURL string
|
||||
}
|
||||
|
||||
func NewMessageQueues(cfg *Options) (map[Backend]MessageQueue, error) {
|
||||
mqs := map[Backend]MessageQueue{}
|
||||
mqs[BackendRedis] = NewRedisMQ()
|
||||
|
||||
if cfg != nil && strings.TrimSpace(cfg.RabbitMQURL) != "" {
|
||||
rmq, err := NewRabbitMQ(RabbitMQConfig{URL: cfg.RabbitMQURL})
|
||||
if err != nil {
|
||||
return mqs, err
|
||||
}
|
||||
mqs[BackendRabbitMQ] = rmq
|
||||
}
|
||||
|
||||
return mqs, nil
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package mq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/pkg/errs"
|
||||
"github.com/rabbitmq/amqp091-go"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type RabbitMQ struct {
|
||||
conn *amqp091.Connection
|
||||
channel *amqp091.Channel
|
||||
}
|
||||
|
||||
type RabbitMQConfig struct {
|
||||
URL string
|
||||
}
|
||||
|
||||
func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
|
||||
conn, err := amqp091.Dial(config.URL)
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(errs.ErrRabbitMQConnectFailed, err)
|
||||
}
|
||||
|
||||
channel, err := conn.Channel()
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
return nil, errs.Wrap(errs.ErrRabbitMQOpenChannelFailed, err)
|
||||
}
|
||||
|
||||
return &RabbitMQ{
|
||||
conn: conn,
|
||||
channel: channel,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (rmq *RabbitMQ) Publish(ctx context.Context, msg PublishMessage) error {
|
||||
topic := strings.TrimSpace(msg.Topic)
|
||||
if topic == "" {
|
||||
return errs.ErrTopicEmpty
|
||||
}
|
||||
if err := rmq.declareTopicExchange(topic); err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQDeclareExchange, err)
|
||||
}
|
||||
|
||||
if err := rmq.channel.PublishWithContext(
|
||||
ctx,
|
||||
topic,
|
||||
"",
|
||||
false,
|
||||
false,
|
||||
amqp091.Publishing{
|
||||
ContentType: "application/json",
|
||||
Body: msg.Body,
|
||||
MessageId: strings.TrimSpace(msg.Key),
|
||||
Headers: toAMQPHeaders(msg.Headers),
|
||||
Timestamp: time.Now(),
|
||||
},
|
||||
); err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQPublishMessage, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (rmq *RabbitMQ) Consume(ctx context.Context, opts ConsumeOptions, handler MessageHandler) error {
|
||||
opts = opts.normalize()
|
||||
topic := strings.TrimSpace(opts.Topic)
|
||||
if topic == "" {
|
||||
return errs.ErrTopicEmpty
|
||||
}
|
||||
|
||||
if err := rmq.declareTopicExchange(topic); err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQDeclareExchange, err)
|
||||
}
|
||||
|
||||
queueName, durable, autoDelete, exclusive := consumeQueueConfig(topic, opts)
|
||||
queue, err := rmq.channel.QueueDeclare(
|
||||
queueName,
|
||||
durable,
|
||||
autoDelete,
|
||||
exclusive,
|
||||
false,
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQDeclareQueue, err)
|
||||
}
|
||||
|
||||
if err := rmq.channel.QueueBind(
|
||||
queue.Name,
|
||||
"",
|
||||
topic,
|
||||
false,
|
||||
nil,
|
||||
); err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQBindQueue, err)
|
||||
}
|
||||
|
||||
consumerTag := buildConsumerTag(opts.Consumer, topic)
|
||||
msgs, err := rmq.channel.Consume(
|
||||
queue.Name,
|
||||
consumerTag,
|
||||
opts.AutoAck,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
return errs.Wrap(errs.ErrRabbitMQRegisterConsumer, err)
|
||||
}
|
||||
|
||||
logrus.Infof("RabbitMQ consume started, topic=%s, queue=%s, mode=%s, group=%s, auto_ack=%v",
|
||||
topic, queue.Name, opts.Mode, opts.Group, opts.AutoAck,
|
||||
)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case d, ok := <-msgs:
|
||||
if !ok {
|
||||
return errs.Detailf(errs.ErrRabbitMQConsumerChanClosed, "topic=%s, queue=%s", topic, queue.Name)
|
||||
}
|
||||
|
||||
message := Message{
|
||||
ID: d.MessageId,
|
||||
Topic: topic,
|
||||
Key: d.MessageId,
|
||||
Body: d.Body,
|
||||
Headers: fromAMQPHeaders(d.Headers),
|
||||
Timestamp: d.Timestamp,
|
||||
}
|
||||
if message.Timestamp.IsZero() {
|
||||
message.Timestamp = time.Now()
|
||||
}
|
||||
|
||||
handleErr := handler(context.Background(), message)
|
||||
if opts.AutoAck {
|
||||
if handleErr != nil {
|
||||
logrus.Errorf("RabbitMQ consume handler failed (auto-ack), topic=%s, queue=%s, err=%v", topic, queue.Name, handleErr)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if handleErr != nil {
|
||||
requeue := opts.OnError == ErrorActionRequeue
|
||||
if nackErr := d.Nack(false, requeue); nackErr != nil {
|
||||
logrus.Errorf("RabbitMQ nack failed, topic=%s, queue=%s, requeue=%v, err=%v", topic, queue.Name, requeue, nackErr)
|
||||
}
|
||||
logrus.Errorf("RabbitMQ consume handler failed, topic=%s, queue=%s, requeue=%v, err=%v", topic, queue.Name, requeue, handleErr)
|
||||
continue
|
||||
}
|
||||
|
||||
if ackErr := d.Ack(false); ackErr != nil {
|
||||
logrus.Errorf("RabbitMQ ack failed, topic=%s, queue=%s, err=%v", topic, queue.Name, ackErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (rmq *RabbitMQ) declareTopicExchange(topic string) error {
|
||||
return rmq.channel.ExchangeDeclare(
|
||||
topic,
|
||||
"topic",
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
nil,
|
||||
)
|
||||
}
|
||||
|
||||
func consumeQueueConfig(topic string, opts ConsumeOptions) (queueName string, durable, autoDelete, exclusive bool) {
|
||||
if opts.Mode != ConsumeModeGroup {
|
||||
return "", false, true, true
|
||||
}
|
||||
group := strings.TrimSpace(opts.Group)
|
||||
if group == "" {
|
||||
group = "default"
|
||||
}
|
||||
return fmt.Sprintf("%s.%s", topic, group), true, false, false
|
||||
}
|
||||
|
||||
func buildConsumerTag(consumer, topic string) string {
|
||||
consumer = strings.TrimSpace(consumer)
|
||||
if consumer == "" {
|
||||
hostName, _ := os.Hostname()
|
||||
consumer = fmt.Sprintf("%s-%s-%d", strings.TrimSpace(topic), strings.TrimSpace(hostName), os.Getpid())
|
||||
}
|
||||
return consumer
|
||||
}
|
||||
|
||||
func toAMQPHeaders(headers map[string]string) amqp091.Table {
|
||||
if len(headers) == 0 {
|
||||
return nil
|
||||
}
|
||||
table := amqp091.Table{}
|
||||
for key, val := range headers {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
table[key] = val
|
||||
}
|
||||
if len(table) == 0 {
|
||||
return nil
|
||||
}
|
||||
return table
|
||||
}
|
||||
|
||||
func fromAMQPHeaders(table amqp091.Table) map[string]string {
|
||||
if len(table) == 0 {
|
||||
return nil
|
||||
}
|
||||
headers := make(map[string]string, len(table))
|
||||
for key, val := range table {
|
||||
headers[key] = fmt.Sprint(val)
|
||||
}
|
||||
return headers
|
||||
}
|
||||
|
||||
func (rmq *RabbitMQ) Close() error {
|
||||
if rmq.channel != nil {
|
||||
rmq.channel.Close()
|
||||
}
|
||||
if rmq.conn != nil {
|
||||
rmq.conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package mq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/pkg/errs"
|
||||
"github.com/XingfenD/yoresee_doc/pkg/storage"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type RedisMQ struct {
|
||||
client *redis.Client
|
||||
}
|
||||
|
||||
func NewRedisMQ() *RedisMQ {
|
||||
return &RedisMQ{
|
||||
client: storage.GetRedis(),
|
||||
}
|
||||
}
|
||||
|
||||
func (rmq *RedisMQ) Publish(ctx context.Context, msg PublishMessage) error {
|
||||
topic := strings.TrimSpace(msg.Topic)
|
||||
if topic == "" {
|
||||
return errs.ErrTopicEmpty
|
||||
}
|
||||
if rmq.client == nil {
|
||||
return errs.ErrRedisClientNotInitialized
|
||||
}
|
||||
return rmq.client.Publish(ctx, topic, msg.Body).Err()
|
||||
}
|
||||
|
||||
func (rmq *RedisMQ) Consume(ctx context.Context, opts ConsumeOptions, handler MessageHandler) error {
|
||||
opts = opts.normalize()
|
||||
topic := strings.TrimSpace(opts.Topic)
|
||||
if topic == "" {
|
||||
return errs.ErrTopicEmpty
|
||||
}
|
||||
if rmq.client == nil {
|
||||
return errs.ErrRedisClientNotInitialized
|
||||
}
|
||||
if opts.Mode == ConsumeModeGroup && strings.TrimSpace(opts.Group) != "" {
|
||||
logrus.Warnf("Redis Pub/Sub does not support true group consume, fallback to fanout mode, topic=%s, group=%s", topic, opts.Group)
|
||||
}
|
||||
|
||||
pubsub := rmq.client.Subscribe(ctx, topic)
|
||||
defer pubsub.Close()
|
||||
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if err := handler(context.Background(), Message{
|
||||
Topic: topic,
|
||||
Body: []byte(msg.Payload),
|
||||
Timestamp: time.Now(),
|
||||
}); err != nil {
|
||||
logrus.Errorf("Redis consume handler failed, topic=%s, err=%v", topic, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (rmq *RedisMQ) Close() error {
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user