从零实现一个消息队列：当一条消息需要被多个不同消费者消费
半年多以前写过一篇《300 行代码实现一个消息队列》，实现了一个最简单的消息队列，能够把消息事件放到异步去处理。但随着业务的发展，简单的异步处理已经不能满足需求了：一条消息需要被多个不同的消费者消费。
按照原先的消息队列设计，要实现这个需求，要么是多注册几个队列，生产者往多个队列发消息；要么是在消费者那多开几个协程去分别处理。但这两种方案都高耦合，很不优雅。
优雅的方案应该是引入一个负责转发消息的中间层，生产者只需要将消息发给这个中间层，再由中间层转发给不同的队列（消费者），这个中间层我们称为 Exchange
队列 Queue 调整
在开始介绍 Exchange 之前，我们需要对 Queue 进行调整。
先前我们定义了一个 type Q struct，它的作用是将 Queue、ComsumerFunc、并发数和消息序号封装在一起，现在我们可以将消费者函数和消费者并发数直接作为 Queue 的属性，并将消息序号交由 Exchange 管理。因此 Q 没有存在的必要了，同时 Queue 修改为：
type Queue struct {
wg sync.WaitGroup
in chan *Msg
out chan *Msg
msgs []*Msg
// 消费者函数
consumer ConsumerFunc
// 消费者并发数
concurrency int
}
// 同时更新实例化函数
func NewQueue(inBufSize, outBufSize, concurrency int, consumer ConsumerFunc) *Queue {
if inBufSize <= 0 {
inBufSize = 1
}
if outBufSize <= 0 {
outBufSize = 1
}
return &Queue{
in: make(chan *Msg, inBufSize),
out: make(chan *Msg, outBufSize),
msgs: make([]*Msg, 0, 1024),
consumer: consumer,
concurrency: concurrency,
}
}然后再新加两个方法，以便做消费者调度：
func (q *Queue) GetConsumer() ConsumerFunc {
return q.consumer
}
func (q *Queue) GetConcurrency() int {
return q.concurrency
}至此，Queue 的改造就完成了。
Exchange
Exchange 应该实现以下接口：
type ExchangeInterface interface {
Publish(any) error // 发布消息
Bind(*Queue) // 注册队列
GetQueues() []*Queue // 获取全部队列（供 MQ 调度）
Start() // 启动 Exchange 的消息转发
Close() // 关闭 Exchange
}具体实现如下：
import (
"sync"
)
type Exchange struct {
mu sync.Mutex
// 绑定的队列列表
queues []*Queue
// 当前消息序号
seq uint64
}
func NewExchange() Exchange {
return &Exchange{}
}
func (e *Exchange) Publish(data any) error {
e.mu.Lock()
defer e.mu.Unlock()
if len(e.queues) == 0 {
return nil // 没有绑定的队列，直接丢弃消息
}
msg := &Msg{
ID: e.seq,
Data: data,
}
e.seq++
for _, q := range e.queues {
cloned := *msg
q.Enqueue(&cloned) // TODO: 可按需增加错误处理逻辑，例如记录失败的消息等
}
return nil
}
func (e *Exchange) Bind(q *Queue) {
e.mu.Lock()
defer e.mu.Unlock()
e.queues = append(e.queues, q)
}
func (e *Exchange) Type() ExchangeType {
return Fanout
}
func (e *Exchange) GetQueues() []*Queue {
return e.queues
}
func (e *Exchange) Start() {
for _, q := range e.queues {
q.Start()
}
}
func (e *Exchange) Close() {
for _, q := range e.queues {
q.Close()
}
}MQ 调整
MQ 最主要的调整就是 Q 废弃后需要将直接调度 queues 改为 exchanges
type MQ struct {
wg sync.WaitGroup
mu sync.Mutex
running bool
- queues map[string]*Q
+ exchanges map[string]*Exchange
log contract.Logger
}
func NewMQ(log contract.Logger) *MQ {
return &MQ{
- queues: make(map[string]*Q),
- log: log,
+ exchanges: make(map[string]Exchange),
+ log: log,
}
}然后修改注册队列的方法：
// 注册一个队列，自动创建对应的 Exchange
func (m *MQ) RegisterQueue(name string, typ ExchangeType, consumer ConsumerFunc, concurrency int) error {
m.mu.Lock()
defer m.mu.Unlock()
if m.running {
return errors.New("mq: cannot register queue while MQ is running")
}
if _, exists := m.exchanges[name]; !exists {
m.exchanges[name] = NewExchange()
}
m.exchanges[name].Bind(NewQueue(concurrency, concurrency, concurrency, consumer))
return nil
}发布消息的方法也需要移除消息序号的逻辑，同时调整为通过 Exchange 去发布消息：
func (m *MQ) Publish(queueName string, data any) error {
m.mu.Lock()
e, exists := m.exchanges[queueName]
m.mu.Unlock()
if !exists {
return errors.New("mq: queue not found")
}
return e.Publish(data)
}相应的，启动消息队列和消费调度的方法也需要调整：
func (m *MQ) Start(ctx context.Context) error {
m.mu.Lock()
defer m.mu.Unlock()
if m.running {
return errors.New("mq: MQ is already running")
}
m.running = true
for exchangeName, exchange := range m.exchanges {
exchange.Start()
for idx, queue := range exchange.GetQueues() {
for workerID := 0; workerID < queue.GetConcurrency(); workerID++ {
m.wg.Add(1)
go m.consume(ctx, exchangeName, workerID, queue)
}
m.log.Infof("Started Exchange[%s].queue[%d] with concurrency %d", exchangeName, idx, queue.GetConcurrency())
}
}
return nil
}
func (m *MQ) consume(ctx context.Context, queueName string, workerID int, q *Queue) {
defer m.wg.Done()
for {
defer func() {
if r := recover(); r != nil {
m.log.Errorf("Recovered from panic in queue[%s] worker[%d]: %v", queueName, workerID, r)
}
}()
select {
case <-ctx.Done():
m.log.Infof("Context canceled. Stopping consumer for queue[%s] worker[%d]", queueName, workerID)
q.Close()
case msg, ok := <-q.Dequeue():
if !ok {
// 队列已关闭，退出
m.log.Infof("Queue[%s] worker[%d] is stopping as the queue is closed", queueName, workerID)
return
}
if err := q.GetConsumer()(ctx, msg); err != nil {
// 处理消息失败，记录日志或进行其他处理
m.log.Errorf("Failed to process message ID[%d] from queue[%s] worker[%d]: %v", msg.ID, queueName, workerID, err)
}
}
}
}最后将停止 MQ 的方法也做相应调整：
func (m *MQ) Stop() {
m.mu.Lock()
if !m.running {
m.mu.Unlock()
return
}
m.running = false
for _, e := range m.exchanges {
e.Close()
}
m.mu.Unlock()
m.wg.Wait()
}至此，我们的消息队列支持同一条消息给多个不同的消费者消费了。没错，细心的你可能发现了，现在这个队列跟 RabbitMQ 的 FanoutExchange 非常像😉
