《游戏服务端编程实践》10.3.1 消息通道与订阅模型

// file: ch10/hub/pubsub.go
// tier: L1
// verify: go test ./ch10/hub/

package hub

import (
	"errors"
	"fmt"
	"strings"
	"sync"
)

var (
	ErrEmptyChannel = errors.New("hub: empty channel")
	ErrBadChannel   = errors.New("hub: invalid channel name")
)

const maxChannelLen = 128

// ValidateChannel 只接受分段式 ASCII 频道名，避免命令注入和空段。
func ValidateChannel(name string) error {
	if name == "" {
		return ErrEmptyChannel
	}
	if len(name) > maxChannelLen {
		return fmt.Errorf("%w: too long", ErrBadChannel)
	}
	for _, part := range strings.Split(name, ":") {
		if part == "" {
			return fmt.Errorf("%w: empty segment", ErrBadChannel)
		}
		for i := 0; i < len(part); i++ {
			if !validByte(part[i]) {
				return fmt.Errorf("%w: byte %q", ErrBadChannel, part[i])
			}
		}
	}
	return nil
}

func validByte(c byte) bool {
	return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' ||
		c == '_' || c == '-' || c == '.'
}

// Message 是 Broker 投递的消息；Pub/Sub 不保存历史，Seq 只用于本地排障。
type Message struct {
	Channel string
	Seq     uint64
	Payload []byte
}

// Broker 模拟 Redis Pub/Sub 的至多一次投递：订阅者不在线就收不到。
type Broker struct {
	mu     sync.Mutex
	nextID uint64
	seq    uint64
	buffer int
	subs   map[string]map[uint64]*Subscription
}

// NewBroker 的 buffer 是每个订阅者的缓冲深度，非正值按 16 处理。
func NewBroker(buffer int) *Broker {
	if buffer <= 0 {
		buffer = 16
	}
	return &Broker{buffer: buffer, subs: make(map[string]map[uint64]*Subscription)}
}

// Subscription 是一条独立订阅，关闭后消息通道随之关闭。
type Subscription struct {
	id       uint64
	channel  string
	broker   *Broker
	messages chan Message
	once     sync.Once
}

func (s *Subscription) Channel() string          { return s.channel }
func (s *Subscription) Messages() <-chan Message { return s.messages }

// Close 幂等；重复关闭不会重复关闭通道。
func (s *Subscription) Close() { s.once.Do(func() { s.broker.unsubscribe(s) }) }

// Subscribe 校验频道名后注册订阅者。
func (b *Broker) Subscribe(channel string) (*Subscription, error) {
	if err := ValidateChannel(channel); err != nil {
		return nil, err
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	b.nextID++
	s := &Subscription{
		id:       b.nextID,
		channel:  channel,
		broker:   b,
		messages: make(chan Message, b.buffer),
	}
	if b.subs[channel] == nil {
		b.subs[channel] = make(map[uint64]*Subscription)
	}
	b.subs[channel][s.id] = s
	return s, nil
}

func (b *Broker) unsubscribe(s *Subscription) {
	b.mu.Lock()
	defer b.mu.Unlock()
	subs, ok := b.subs[s.channel]
	if !ok {
		return
	}
	if _, ok := subs[s.id]; !ok {
		return
	}
	delete(subs, s.id)
	if len(subs) == 0 {
		delete(b.subs, s.channel)
	}
	close(s.messages)
}

// Publish 返回入队的订阅者数量；缓冲满的订阅者被跳过，发布者不会被阻塞。
func (b *Broker) Publish(channel string, payload []byte) (int, error) {
	if err := ValidateChannel(channel); err != nil {
		return 0, err
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	b.seq++
	msg := Message{Channel: channel, Seq: b.seq, Payload: append([]byte(nil), payload...)}
	delivered := 0
	for _, s := range b.subs[channel] {
		select {
		case s.messages <- msg:
			delivered++
		default:
		}
	}
	return delivered, nil
}

// Session 表示一条客户端订阅连接：断线时订阅全部失效，重连后按原频道重建。
type Session struct {
	broker   *Broker
	channels []string
	mu       sync.Mutex
	subs     []*Subscription
}

// NewSession 去重频道并按给定顺序订阅。
func NewSession(b *Broker, channels ...string) (*Session, error) {
	s := &Session{broker: b}
	seen := make(map[string]struct{}, len(channels))
	for _, ch := range channels {
		if err := ValidateChannel(ch); err != nil {
			return nil, err
		}
		if _, ok := seen[ch]; ok {
			continue
		}
		seen[ch] = struct{}{}
		s.channels = append(s.channels, ch)
	}
	return s, s.subscribeAll()
}

func (s *Session) subscribeAll() error {
	subs := make([]*Subscription, 0, len(s.channels))
	for _, ch := range s.channels {
		sub, err := s.broker.Subscribe(ch)
		if err != nil {
			for _, old := range subs {
				old.Close()
			}
			return err
		}
		subs = append(subs, sub)
	}
	s.subs = subs
	return nil
}

func (s *Session) Subscriptions() []*Subscription {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]*Subscription(nil), s.subs...)
}

// Disconnect 模拟断线：订阅立即失效，期间的消息不会补发。
func (s *Session) Disconnect() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sub := range s.subs {
		sub.Close()
	}
	s.subs = nil
}

// Reconnect 关闭残留订阅后按原频道重建，重复调用不会产生重复订阅。
func (s *Session) Reconnect() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, sub := range s.subs {
		sub.Close()
	}
	s.subs = nil
	return s.subscribeAll()
}

cd .local/checks/gs-practice-lab && gofmt -l ch10/hub && go vet ./ch10/hub/ && go test ./ch10/hub/ -count=1

ok  	gspractice/ch10/hub	0.275s