AutomatedTesting/usecases_server/utils/pubsub/pubsub.go

package pubsub

import (
	"strings"
	"sync"
	"time"
)

type Message struct {
	Topic   string
	Payload []byte
}

type (
	subscriber  chan Message      // 订阅者为一个管道
	topicFilter func(string) bool // 主题为一个过滤器
)

// 发布者对象
type Publisher struct {
	m           sync.RWMutex               // 读写锁
	buffer      int                        // 订阅队列的缓存大小
	timeout     time.Duration              // 发布超时时间
	subscribers map[subscriber]topicFilter // 所有订阅者信息
}

// 构建一个发布者对象，可以设置发布超时时间和缓存队列的长度
func NewPublisher(publishTimeout time.Duration, buffer int) *Publisher {
	return &Publisher{
		buffer:      buffer,
		timeout:     publishTimeout,
		subscribers: make(map[subscriber]topicFilter),
	}
}

// 添加一个新的订阅者，订阅过滤筛选后的主题
func (p *Publisher) subscribeTopic(filter topicFilter) chan Message {
	ch := make(chan Message, p.buffer)
	p.m.Lock()
	p.subscribers[ch] = filter
	p.m.Unlock()
	return ch
}

func (p *Publisher) Subscribe(topic string) chan Message {
	var messages chan Message
	if topic != "" {
		if strings.Contains(topic, "*") {
			sps := strings.Split(topic, "*")
			if len(sps) >= 2 && len(sps[0]) > 0 && len(sps[1]) >= 0 {
				messages = p.subscribeTopic(func(t string) bool {
					flag := true
					for _, sp := range sps {
						if !strings.Contains(t, sp) {
							flag = false
							break
						}
					}
					return flag
				})
			} else {
				messages = p.subscribeTopic(nil)
			}
		} else {
			messages = p.subscribeTopic(func(t string) bool {
				if t == topic {
					return true
				}
				return false
			})
		}
	} else {
		messages = p.subscribeTopic(nil)
	}
	return messages
}

// 退出订阅
func (p *Publisher) Unsubscribe(sub chan Message) {
	p.m.Lock()
	defer p.m.Unlock()
	delete(p.subscribers, sub)
	close(sub)
}

// 发送主题，可以容忍一定的超时
func (p *Publisher) sendTopic(sub subscriber, filter topicFilter, msg Message, wg *sync.WaitGroup) {
	defer wg.Done()
	// 没有订阅该主题
	if filter != nil && !filter(msg.Topic) {
		return
	}
	select {
	case sub <- msg:
	case <-time.After(p.timeout):
	}
}

// 发布一个主题
func (p *Publisher) Publish(topic string, payload []byte) {
	p.m.RLock()
	defer p.m.RUnlock()
	var wg sync.WaitGroup
	msg := Message{
		Topic:   topic,
		Payload: payload,
	}
	for sub, filter := range p.subscribers {
		wg.Add(1)
		go p.sendTopic(sub, filter, msg, &wg)
	}
	wg.Wait()
}

// 关闭发布者对象，同时关闭所有的订阅者管道
func (p *Publisher) Close() {
	p.m.Lock()
	defer p.m.Unlock()
	for sub := range p.subscribers {
		delete(p.subscribers, sub)
		close(sub)
	}
}

var PubSub *Publisher

func init() {
	PubSub = NewPublisher(100*time.Millisecond, 1)
}