sig-pub/internal/exchange/publisher.go

package exchange

import "sig-pub/pkg/utils/collect"

type Publisher[TID comparable, T any] struct {
	m0 *collect.ConcurrentMap[string, map[TID]T]
}

func NewPublisher[TID comparable, T any](concurrentLevel int) *Publisher[TID, T] {
	m := collect.NewConcurrentMap[string, map[TID]T](concurrentLevel, func(k string) string { return k })
	return &Publisher[TID, T]{
		m0: m,
	}
}

// Subscribe 订阅
func (p *Publisher[TID, T]) Subscribe(k string, tid TID, ele T) {
	_ = p.m0.LoadAndUpdate(k, func(v map[TID]T) (remove bool, nextV map[TID]T) {
		if v == nil {
			v = make(map[TID]T)
		}
		v[tid] = ele
		return false, v
	})
}

// Unsubscribe 取消订阅
func (p *Publisher[TID, T]) Unsubscribe(k string, tid TID) {
	p.m0.LoadAndUpdate(k, func(v map[TID]T) (remove bool, nextV map[TID]T) {
		delete(v, tid)
		if len(v) == 0 {
			return true, v
		}
		return false, v
	})
}

// UnsubscribeAll 取消所有订阅
func (p *Publisher[TID, T]) UnsubscribeAll(tid TID) {
	p.m0.RangeUpdate(func(key string, v map[TID]T) (next bool, remove bool, newV map[TID]T) {
		delete(v, tid)
		remove = len(v) == 0
		return true, remove, v
	})
}

// Publisher 获取匹配的 subscribers
func (p *Publisher[TID, T]) Publisher(k string) (subs []T) {
	p.m0.LoadRLock(k, func(subIds map[TID]T, ok bool) {
		if !ok {
			return
		}
		for _, sub := range subIds {
			subs = append(subs, sub)
		}
	})
	// todo cache subs
	return
}