what
1322280daf
- 新增 pending 缓冲区，publish 时若无订阅者则暂存消息
- subscribe 时自动将缓冲消息投入 channel，解决服务重启后恢复任务丢失的问题
- 去除 broadcast 5ms 超时导致的消息丢失
- chan bool 改为 chan struct{}，RWMutex 改为 Mutex
- 新增 broker_test.go，12 个单元测试覆盖核心场景（含 -race）
- 为 client_test.go 中的无限循环 demo 添加 t.Skip()
121 lines
2.3 KiB
Go
121 lines
2.3 KiB
Go
package queue
|
||
|
||
import (
|
||
"errors"
|
||
"sync"
|
||
)
|
||
|
||
type Broker struct {
|
||
exit chan struct{}
|
||
capacity int
|
||
topics map[string][]chan any
|
||
pending map[string][]any // 订阅前发布的消息缓冲，subscribe 时一次性投递
|
||
mu sync.Mutex
|
||
}
|
||
|
||
func (b *Broker) setConditions(capacity int) {
|
||
b.mu.Lock()
|
||
b.capacity = capacity
|
||
b.mu.Unlock()
|
||
}
|
||
|
||
func (b *Broker) close() {
|
||
select {
|
||
case <-b.exit:
|
||
return
|
||
default:
|
||
close(b.exit)
|
||
b.mu.Lock()
|
||
b.topics = make(map[string][]chan any)
|
||
b.pending = make(map[string][]any)
|
||
b.mu.Unlock()
|
||
}
|
||
}
|
||
|
||
// publish 推送消息；若暂无订阅者则缓冲，等待订阅者注册后投递。
|
||
func (b *Broker) publish(topic string, msg any) error {
|
||
select {
|
||
case <-b.exit:
|
||
return errors.New("broker closed")
|
||
default:
|
||
}
|
||
|
||
b.mu.Lock()
|
||
subs := b.topics[topic]
|
||
if len(subs) == 0 {
|
||
b.pending[topic] = append(b.pending[topic], msg)
|
||
b.mu.Unlock()
|
||
return nil
|
||
}
|
||
// 持有锁期间只做列表复制，发送在锁外进行，避免阻塞其他 publish
|
||
chs := make([]chan any, len(subs))
|
||
copy(chs, subs)
|
||
b.mu.Unlock()
|
||
|
||
for _, ch := range chs {
|
||
select {
|
||
case ch <- msg:
|
||
case <-b.exit:
|
||
return errors.New("broker closed")
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// subscribe 订阅 topic，返回 channel；同时将该 topic 的缓冲消息立即投入 channel。
|
||
func (b *Broker) subscribe(topic string) (<-chan any, error) {
|
||
select {
|
||
case <-b.exit:
|
||
return nil, errors.New("broker closed")
|
||
default:
|
||
}
|
||
|
||
b.mu.Lock()
|
||
capacity := b.capacity
|
||
if capacity <= 0 {
|
||
capacity = 10
|
||
}
|
||
ch := make(chan any, capacity)
|
||
b.topics[topic] = append(b.topics[topic], ch)
|
||
buffered := b.pending[topic]
|
||
delete(b.pending, topic)
|
||
b.mu.Unlock()
|
||
|
||
// channel 刚创建必然不满，直接写入不会阻塞
|
||
for _, msg := range buffered {
|
||
ch <- msg
|
||
}
|
||
|
||
return ch, nil
|
||
}
|
||
|
||
func (b *Broker) unsubscribe(topic string, sub <-chan any) error {
|
||
select {
|
||
case <-b.exit:
|
||
return errors.New("broker closed")
|
||
default:
|
||
}
|
||
|
||
b.mu.Lock()
|
||
defer b.mu.Unlock()
|
||
|
||
subs := b.topics[topic]
|
||
newSubs := subs[:0]
|
||
for _, s := range subs {
|
||
if s != sub {
|
||
newSubs = append(newSubs, s)
|
||
}
|
||
}
|
||
b.topics[topic] = newSubs
|
||
return nil
|
||
}
|
||
|
||
func NewBroker() *Broker {
|
||
return &Broker{
|
||
exit: make(chan struct{}),
|
||
capacity: 10,
|
||
topics: make(map[string][]chan any),
|
||
pending: make(map[string][]any),
|
||
}
|
||
}
|