first commit

broker.go

@@ -0,0 +1,153 @@
+package queue
+
+import (
+	"errors"
+	"sync"
+	"time"
+)
+
+type Broker struct {
+	exit         chan bool             // 关闭消息队列通道
+	capacity     int                   // 消息队列的容量
+	topics       map[string][]chan any // key： topic  value ： queue, 一个topic可以有多个订阅者，一个订阅者对应着一个通道
+	sync.RWMutex                       // 同步锁
+}
+
+// 设置消息容量
+// @description 控制消息队列的大小
+func (b *Broker) setConditions(capacity int) {
+	b.capacity = capacity
+}
+
+// 关闭消息队列
+func (b *Broker) close() {
+	select {
+	case <-b.exit:
+		return
+	default:
+		close(b.exit)
+		b.Lock()
+		b.topics = make(map[string][]chan any)
+		b.Unlock()
+	}
+	return
+}
+
+// 消息推送
+// @param topic 订阅的主题
+// @param msg 传递的消息
+func (b *Broker) publish(topic string, pub any) error {
+	select {
+	case <-b.exit:
+		return errors.New("broker closed")
+	default:
+	}
+
+	b.RLock()
+	subscribers, ok := b.topics[topic]
+	b.RUnlock()
+	if !ok {
+		return nil
+	}
+
+	b.broadcast(pub, subscribers)
+	return nil
+}
+
+// 消息广播
+// @description 对推送的消息进行广播，保证每一个订阅者都可以收到
+func (b *Broker) broadcast(msg any, subscribers []chan any) {
+	count := len(subscribers)
+	concurrency := 1
+
+	switch {
+	case count > 1000:
+		concurrency = 3
+	case count > 100:
+		concurrency = 2
+	default:
+		concurrency = 1
+	}
+
+	pub := func(start int) {
+		//采用Timer 而不是使用time.After 原因：time.After会产生内存泄漏 在计时器触发之前，垃圾回收器不会回收Timer
+		idleDuration := 5 * time.Millisecond
+		idleTimeout := time.NewTimer(idleDuration)
+		defer idleTimeout.Stop()
+		for j := start; j < count; j += concurrency {
+			if !idleTimeout.Stop() {
+				select {
+				case <-idleTimeout.C:
+				default:
+				}
+			}
+			idleTimeout.Reset(idleDuration)
+			select {
+			case subscribers[j] <- msg:
+			case <-idleTimeout.C:
+			case <-b.exit:
+				return
+			}
+		}
+	}
+	for i := 0; i < concurrency; i++ {
+		go pub(i)
+	}
+}
+
+// 消息订阅
+// @description 传入订阅的主题，即可完成订阅
+// @param topic 订阅的主题
+// @return sub 通道用来接收数据
+func (b *Broker) subscribe(topic string) (<-chan any, error) {
+	select {
+	case <-b.exit:
+		return nil, errors.New("broker closed")
+	default:
+	}
+
+	ch := make(chan any, b.capacity)
+	b.Lock()
+	b.topics[topic] = append(b.topics[topic], ch)
+	b.Unlock()
+	return ch, nil
+}
+
+// 取消订阅
+// @param topic 订阅的主题
+// @param sub 消息订阅的通道
+func (b *Broker) unsubscribe(topic string, sub <-chan any) error {
+	select {
+	case <-b.exit:
+		return errors.New("broker closed")
+	default:
+	}
+
+	b.RLock()
+	subscribers, ok := b.topics[topic]
+	b.RUnlock()
+
+	if !ok {
+		return nil
+	}
+	// delete subscriber
+	b.Lock()
+	var newSubs []chan any
+	for _, subscriber := range subscribers {
+		if subscriber == sub {
+			continue
+		}
+		newSubs = append(newSubs, subscriber)
+	}
+
+	b.topics[topic] = newSubs
+	b.Unlock()
+	return nil
+}
+
+func NewBroker() *Broker {
+	return &Broker{
+		exit:   make(chan bool),
+		topics: make(map[string][]chan any),
+	}
+}