// @tag pubsub // @tag core package pubsub import ( "fmt" "regexp" "slices" "strings" "sync" ) // Hub maintains the mapping between channels and their subscribers. type Hub struct { mu sync.RWMutex subs map[string]map[*Client]struct{} patterns map[string]*patternEntry // Pattern index for fast channel-to-pattern matching // Maps channel name to patterns that would match it patternIndex map[string][]string // Client registry for CLIENT LIST/INFO/KILL commands clients map[uint64]*Client } // patternEntry holds compiled regex and subscribers for a pattern. type patternEntry struct { regex *regexp.Regexp clients map[*Client]struct{} } // patternCache caches compiled regex patterns to avoid recompilation. type patternCacheType struct { sync.Map } var patternCache = patternCacheType{} // NewHub creates a new pubsub hub instance ready for use. func NewHub() *Hub { return &Hub{ subs: make(map[string]map[*Client]struct{}), patterns: make(map[string]*patternEntry), patternIndex: make(map[string][]string), clients: make(map[uint64]*Client), } } // RegisterClient adds a client to the registry func (h *Hub) RegisterClient(client *Client) { h.mu.Lock() defer h.mu.Unlock() h.clients[client.ID()] = client } // UnregisterClient removes a client from the registry func (h *Hub) UnregisterClient(client *Client) { h.mu.Lock() defer h.mu.Unlock() delete(h.clients, client.ID()) } // GetClient returns a client by ID, or nil if not found func (h *Hub) GetClient(id uint64) *Client { h.mu.RLock() defer h.mu.RUnlock() return h.clients[id] } // GetAllClients returns a slice of all registered clients func (h *Hub) GetAllClients() []*Client { h.mu.RLock() defer h.mu.RUnlock() clients := make([]*Client, 0, len(h.clients)) for _, c := range h.clients { clients = append(clients, c) } return clients } // ClientCount returns the number of connected clients func (h *Hub) ClientCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients) } // Subscribe adds a client to a channel. func (h *Hub) Subscribe(client *Client, channel string) bool { if channel == "" { return false } // Must hold hub mutex when modifying the global subs map h.mu.Lock() defer h.mu.Unlock() clients, exists := h.subs[channel] if !exists { clients = make(map[*Client]struct{}) h.subs[channel] = clients // New channel - add to pattern index h.addChannelToPatternIndex(channel) } if _, already := clients[client]; already { return false } clients[client] = struct{}{} client.AddSubscription(channel) return true } // addChannelToPatternIndex adds a channel to the pattern index for all matching patterns. func (h *Hub) addChannelToPatternIndex(channel string) { for pattern, entry := range h.patterns { if entry.regex.MatchString(channel) { h.patternIndex[channel] = append(h.patternIndex[channel], pattern) } } } // Unsubscribe removes a client from a channel. func (h *Hub) Unsubscribe(client *Client, channel string) bool { if channel == "" { return false } // Must hold hub mutex when modifying the global subs map h.mu.Lock() defer h.mu.Unlock() clients, exists := h.subs[channel] if !exists { return false } if _, ok := clients[client]; !ok { return false } delete(clients, client) if len(clients) == 0 { delete(h.subs, channel) } client.RemoveSubscription(channel) return true } // UnsubscribeAll removes a client from all subscribed channels and patterns. func (h *Hub) UnsubscribeAll(client *Client) { channels := client.GetChannels() for _, channel := range channels { h.Unsubscribe(client, channel) } patterns := client.GetPatterns() for _, pattern := range patterns { h.PatternUnsubscribe(client, pattern) } } // Publish sends a message to all subscribers of a channel. func (h *Hub) Publish(channel string, message []byte) int { // Copy clients under lock, then release lock before sending h.mu.RLock() clients, exists := h.subs[channel] if !exists || len(clients) == 0 { h.mu.RUnlock() return 0 } clientList := make([]*Client, 0, len(clients)) for client := range clients { clientList = append(clientList, client) } h.mu.RUnlock() count := 0 for _, client := range clientList { if err := client.sendMessage(channel, message); err != nil { h.removeClientFromChannel(client, channel) continue } count++ } return count } // PublishToPattern sends a message to all clients subscribed to patterns matching the channel. // nolint: cyclop func (h *Hub) PublishToPattern(channel string, message []byte) int { h.mu.RLock() // First check the pattern index for patterns that have been indexed // Copy the slice to avoid aliasing issues when appending matchedPatterns := slices.Clone(h.patternIndex[channel]) // Also check all patterns to see if any match this channel // This ensures we don't miss patterns when publishing to channels with no direct subscribers for pattern, entry := range h.patterns { if entry.regex.MatchString(channel) { // Check if already in matchedPatterns to avoid duplicates found := false for _, p := range matchedPatterns { if p == pattern { found = true break } } if !found { matchedPatterns = append(matchedPatterns, pattern) } } } if len(matchedPatterns) == 0 { h.mu.RUnlock() return 0 } type match struct { client *Client pattern string } var matches []match for _, pattern := range matchedPatterns { if entry, ok := h.patterns[pattern]; ok { for client := range entry.clients { matches = append(matches, match{client: client, pattern: pattern}) } } } h.mu.RUnlock() count := 0 for _, m := range matches { if err := m.client.sendPMessage(m.pattern, channel, message); err != nil { h.removeClientFromPattern(m.client, m.pattern) continue } count++ } return count } // PatternSubscribe adds a client to a pattern subscription. func (h *Hub) PatternSubscribe(client *Client, pattern string) (bool, error) { regex, err := getOrCompilePattern(pattern) if err != nil { return false, fmt.Errorf("invalid pattern: %v", err) } h.mu.Lock() defer h.mu.Unlock() entry, exists := h.patterns[pattern] if !exists { entry = &patternEntry{ regex: regex, clients: make(map[*Client]struct{}), } h.patterns[pattern] = entry } if _, already := entry.clients[client]; already { return false, nil } entry.clients[client] = struct{}{} // Add pattern to index for all existing channels for ch := range h.subs { if regex.MatchString(ch) { h.patternIndex[ch] = append(h.patternIndex[ch], pattern) } } client.AddPatternSubscription(pattern) return true, nil } // getOrCompilePattern returns a cached compiled regex or compiles and caches a new one. func getOrCompilePattern(pattern string) (*regexp.Regexp, error) { if cached, ok := patternCache.Load(pattern); ok { return cached.(*regexp.Regexp), nil } regex, err := compilePattern(pattern) if err != nil { return nil, err } patternCache.Store(pattern, regex) return regex, nil } func compilePattern(pattern string) (*regexp.Regexp, error) { var buf strings.Builder for i := 0; i < len(pattern); i++ { switch pattern[i] { case '*': buf.WriteString(".*") case '?': buf.WriteByte('.') case '\\': if i+1 < len(pattern) { buf.WriteByte('\\') buf.WriteByte(pattern[i+1]) i++ } case '.', '+', '(', ')', '[', ']', '{', '}', '^', '$', '|': buf.WriteByte('\\') buf.WriteByte(pattern[i]) default: buf.WriteByte(pattern[i]) } } return regexp.Compile("^" + buf.String() + "$") } // PatternUnsubscribe removes a client from a pattern subscription. func (h *Hub) PatternUnsubscribe(client *Client, pattern string) bool { h.mu.Lock() defer h.mu.Unlock() entry, exists := h.patterns[pattern] if !exists { return false } if _, ok := entry.clients[client]; !ok { return false } delete(entry.clients, client) if len(entry.clients) == 0 { delete(h.patterns, pattern) // Remove pattern from index h.removePatternFromIndex(pattern) } client.RemovePatternSubscription(pattern) return true } // removePatternFromIndex removes a pattern from all channel indexes. func (h *Hub) removePatternFromIndex(pattern string) { for ch, patterns := range h.patternIndex { newPatterns := make([]string, 0, len(patterns)-1) for _, p := range patterns { if p != pattern { newPatterns = append(newPatterns, p) } } if len(newPatterns) == 0 { delete(h.patternIndex, ch) } else { h.patternIndex[ch] = newPatterns } } } // PatternCount returns the total number of pattern subscriptions. func (h *Hub) PatternCount() int { h.mu.RLock() defer h.mu.RUnlock() count := 0 for _, entry := range h.patterns { count += len(entry.clients) } return count } // RemoveClientFromChannel handles cleanup when a client's connection fails. func (h *Hub) removeClientFromChannel(client *Client, channel string) { h.mu.Lock() defer h.mu.Unlock() if clients, ok := h.subs[channel]; ok { if _, exists := clients[client]; exists { delete(clients, client) if len(clients) == 0 { delete(h.subs, channel) } } } client.RemoveSubscription(channel) } // RemoveClientFromPattern handles cleanup when a client's connection fails for patterns. func (h *Hub) removeClientFromPattern(client *Client, pattern string) { h.mu.Lock() defer h.mu.Unlock() if entry, ok := h.patterns[pattern]; ok { if _, exists := entry.clients[client]; exists { delete(entry.clients, client) if len(entry.clients) == 0 { delete(h.patterns, pattern) } } } client.RemovePatternSubscription(pattern) } // ChannelCount returns the number of subscribers to a specific channel. func (h *Hub) ChannelCount(channel string) int { h.mu.RLock() defer h.mu.RUnlock() if clients, ok := h.subs[channel]; ok { return len(clients) } return 0 } // TotalChannels returns the total number of active channels. func (h *Hub) TotalChannels() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.subs) } // ChannelList returns a list of all active channel names, optionally filtered by pattern. func (h *Hub) ChannelList(pattern string) []string { h.mu.RLock() defer h.mu.RUnlock() if pattern == "" { channels := make([]string, 0, len(h.subs)) for ch := range h.subs { channels = append(channels, ch) } return channels } // Compile pattern once outside the loop (using cache) regex, err := getOrCompilePattern(pattern) if err != nil { return nil } channels := make([]string, 0, len(h.subs)) for ch := range h.subs { if regex.MatchString(ch) { channels = append(channels, ch) } } return channels }