Initial commit

2026-04-02 23:56:28 +03:00
commit 9e30243533
34 changed files with 8413 additions and 0 deletions
--- a/internal/pubsub/hub.go
+++ b/internal/pubsub/hub.go
@@ -0,0 +1,463 @@
+// @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
+}