Initial commit

internal/pubsub/client.go

@@ -0,0 +1,257 @@
+// @tag pubsub
+// @tag core
+package pubsub
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/eleksir/pekoe/internal/protocol"
+)
+
+// ReplyMode controls client reply behavior
+type ReplyMode int
+
+const (
+	ReplyOn   ReplyMode = iota // Normal reply mode - send all responses
+	ReplyOff                   // Reply off - suppress all responses
+	ReplySkip                  // Reply skip - skip the next response only
+)
+
+// Global client ID counter - each client gets a unique ID
+var clientIDCounter atomic.Uint64
+
+// Client represents a subscriber connection to the pubsub system.
+// It maintains the connection (net.Conn) and tracks which channels/patterns the client is subscribed to.
+type Client struct {
+	id         uint64 // Unique client ID
+	conn       net.Conn
+	name       string    // Connection name set via CLIENT SETNAME
+	replyMode  ReplyMode // Reply mode (ON/OFF/SKIP)
+	flags      uint64    // Client flags (e.g., paused, no-touch)
+	created    time.Time // When the client connected
+	lastActive time.Time // Last activity time
+	subscribed map[string]struct{}
+	patterns   map[string]struct{}
+	mu         sync.RWMutex
+	writeMu    sync.Mutex // Protects concurrent writes to conn
+}
+
+// NewClient creates a new pubsub client wrapping a connection.
+func NewClient(conn net.Conn) *Client {
+	now := time.Now()
+	return &Client{
+		id:         clientIDCounter.Add(1),
+		conn:       conn,
+		name:       "",
+		replyMode:  ReplyOn,
+		flags:      0,
+		created:    now,
+		lastActive: now,
+		subscribed: make(map[string]struct{}),
+		patterns:   make(map[string]struct{}),
+	}
+}
+
+// ID returns the unique client ID
+func (c *Client) ID() uint64 {
+	return c.id
+}
+
+// SetName sets the connection name
+func (c *Client) SetName(name string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.name = name
+}
+
+// GetName returns the connection name
+func (c *Client) GetName() string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.name
+}
+
+// SetReplyMode sets the reply mode
+func (c *Client) SetReplyMode(mode ReplyMode) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.replyMode = mode
+}
+
+// GetReplyMode returns the current reply mode
+func (c *Client) GetReplyMode() ReplyMode {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.replyMode
+}
+
+// UpdateActivity updates the last activity timestamp
+func (c *Client) UpdateActivity() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.lastActive = time.Now()
+}
+
+// GetChannels returns a copy of the channels this client is subscribed to.
+// Returns a slice to prevent external modification.
+func (c *Client) GetChannels() []string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	channels := make([]string, 0, len(c.subscribed))
+	for ch := range c.subscribed {
+		channels = append(channels, ch)
+	}
+	return channels
+}
+
+// GetPatterns returns a copy of the patterns this client is subscribed to.
+func (c *Client) GetPatterns() []string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	patterns := make([]string, 0, len(c.patterns))
+	for p := range c.patterns {
+		patterns = append(patterns, p)
+	}
+	return patterns
+}
+
+// SubsCount returns the number of channels this client is subscribed to.
+func (c *Client) SubsCount() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return len(c.subscribed)
+}
+
+// PatternSubsCount returns the number of patterns this client is subscribed to.
+func (c *Client) PatternSubsCount() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return len(c.patterns)
+}
+
+// AddSubscription tracks that this client has subscribed to a channel.
+func (c *Client) AddSubscription(channel string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.subscribed[channel] = struct{}{}
+	c.lastActive = time.Now()
+}
+
+// RemoveSubscription removes a channel from the client's subscriptions.
+func (c *Client) RemoveSubscription(channel string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	delete(c.subscribed, channel)
+	c.lastActive = time.Now()
+}
+
+// AddPatternSubscription tracks that this client has subscribed to a pattern.
+func (c *Client) AddPatternSubscription(pattern string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.patterns[pattern] = struct{}{}
+	c.lastActive = time.Now()
+}
+
+// RemovePatternSubscription removes a pattern from the client's subscriptions.
+func (c *Client) RemovePatternSubscription(pattern string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	delete(c.patterns, pattern)
+	c.lastActive = time.Now()
+}
+
+// SendMessage sends a pubsub message to the client in RESP format.
+// Format: ["message", channel, payload]
+// Returns error if the write fails - caller will handle cleanup.
+func (c *Client) sendMessage(channel string, message []byte) error {
+	response := protocol.MakeMultiBulk([]protocol.Value{
+		protocol.Bulk("message"),
+		protocol.Bulk(channel),
+		protocol.Bulk(message),
+	})
+
+	c.writeMu.Lock()
+	_, err := c.conn.Write(response)
+	c.writeMu.Unlock()
+	c.UpdateActivity()
+	return err
+}
+
+// SendPMessage sends a pubsub message matching a pattern.
+// Format: ["pmessage", pattern, channel, payload]
+func (c *Client) sendPMessage(pattern string, channel string, message []byte) error {
+	response := protocol.MakeMultiBulk([]protocol.Value{
+		protocol.Bulk("pmessage"),
+		protocol.Bulk(pattern),
+		protocol.Bulk(channel),
+		protocol.Bulk(message),
+	})
+
+	c.writeMu.Lock()
+	_, err := c.conn.Write(response)
+	c.writeMu.Unlock()
+	c.UpdateActivity()
+	return err
+}
+
+// RemoteAddr returns the remote address of the client for logging.
+func (c *Client) RemoteAddr() string {
+	return c.conn.RemoteAddr().String()
+}
+
+// LocalAddr returns the local address of the client connection.
+func (c *Client) LocalAddr() string {
+	return c.conn.LocalAddr().String()
+}
+
+// InfoString returns CLIENT INFO formatted string (Redis format)
+// Format: id= addr= laddr= fd= age= idle= flags= db= sub= psub= multi= qbuf= qbuf-free= obl= oll= omem= tot-mem= cmd= user= redir=
+func (c *Client) InfoString() string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	age := int(time.Since(c.created).Seconds())
+	idle := int(time.Since(c.lastActive).Seconds())
+
+	// Determine flags
+	flags := ""
+	if c.replyMode == ReplyOff {
+		flags += "O"
+	}
+	if c.replyMode == ReplySkip {
+		flags += "S"
+	}
+
+	// Format: Redis CLIENT INFO format
+	return fmt.Sprintf("id=%d addr=%s laddr=%s fd=-1 age=%d idle=%d flags=%s db=0 sub=%d psub=%d multi=-1 qbuf=0 qbuf-free=0 obl=0 oll=0 omem=0 tot-mem=0 cmd=client user=default redir=-1",
+		c.id,
+		c.RemoteAddr(),
+		c.LocalAddr(),
+		age,
+		idle,
+		flags,
+		len(c.subscribed),
+		len(c.patterns),
+	)
+}
+
+// ListString returns CLIENT LIST formatted string
+// Format: id= addr= ... (same as INFO but one line per client)
+func (c *Client) ListString() string {
+	return c.InfoString()
+}
+
+// Close closes the client's connection.
+// Called when the client disconnects or on error.
+func (c *Client) Close() error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	return c.conn.Close()
+}

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
+}

internal/pubsub/pubsub_test.go

@@ -0,0 +1,333 @@
+package pubsub
+
+import (
+	"net"
+	"testing"
+	"time"
+)
+
+type mockConn struct {
+	data []byte
+	pos  int
+}
+
+func (m *mockConn) Read(b []byte) (n int, err error) {
+	if m.pos >= len(m.data) {
+		return 0, nil
+	}
+	n = copy(b, m.data[m.pos:])
+	m.pos += n
+	return n, nil
+}
+
+func (m *mockConn) Write(b []byte) (n int, err error) {
+	m.data = append(m.data, b...)
+	return len(b), nil
+}
+
+func (m *mockConn) Close() error                     { return nil }
+func (m *mockConn) LocalAddr() net.Addr              { return nil }
+func (m *mockConn) RemoteAddr() net.Addr             { return nil }
+func (m *mockConn) SetDeadline(time.Time) error      { return nil }
+func (m *mockConn) SetReadDeadline(time.Time) error  { return nil }
+func (m *mockConn) SetWriteDeadline(time.Time) error { return nil }
+
+func NewMockClient() *Client {
+	return NewClient(&mockConn{data: []byte{}})
+}
+
+func TestNewHub(t *testing.T) {
+	hub := NewHub()
+	if hub == nil {
+		t.Fatal("NewHub() returned nil")
+	}
+	if hub.subs == nil {
+		t.Error("hub.subs is nil")
+	}
+}
+
+func TestHubSubscribe(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	// Subscribe to a channel
+	ok := hub.Subscribe(client, "news")
+	if !ok {
+		t.Error("Subscribe() returned false for new subscription")
+	}
+
+	// Check subscriber count
+	count := hub.ChannelCount("news")
+	if count != 1 {
+		t.Errorf("ChannelCount() = %d, want 1", count)
+	}
+
+	// Duplicate subscription should return false
+	ok = hub.Subscribe(client, "news")
+	if ok {
+		t.Error("Subscribe() should return false for duplicate subscription")
+	}
+}
+
+func TestHubUnsubscribe(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	hub.Subscribe(client, "news")
+
+	ok := hub.Unsubscribe(client, "news")
+	if !ok {
+		t.Error("Unsubscribe() returned false")
+	}
+
+	count := hub.ChannelCount("news")
+	if count != 0 {
+		t.Errorf("ChannelCount() = %d, want 0", count)
+	}
+}
+
+func TestHubUnsubscribeNotSubscribed(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	ok := hub.Unsubscribe(client, "nonexistent")
+	if ok {
+		t.Error("Unsubscribe() should return false for non-subscribed channel")
+	}
+}
+
+func TestHubUnsubscribeAll(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	hub.Subscribe(client, "news")
+	hub.Subscribe(client, "sports")
+	hub.Subscribe(client, "weather")
+
+	hub.UnsubscribeAll(client)
+
+	if hub.ChannelCount("news") != 0 {
+		t.Error("news channel should be empty")
+	}
+	if hub.ChannelCount("sports") != 0 {
+		t.Error("sports channel should be empty")
+	}
+	if hub.ChannelCount("weather") != 0 {
+		t.Error("weather channel should be empty")
+	}
+}
+
+func TestHubPublish(t *testing.T) {
+	hub := NewHub()
+	client1 := NewMockClient()
+	client2 := NewMockClient()
+
+	hub.Subscribe(client1, "news")
+	hub.Subscribe(client2, "news")
+
+	count := hub.Publish("news", []byte("Hello"))
+	if count != 2 {
+		t.Errorf("Publish() = %d, want 2", count)
+	}
+}
+
+func TestHubPublishEmptyChannel(t *testing.T) {
+	hub := NewHub()
+
+	count := hub.Publish("empty", []byte("Hello"))
+	if count != 0 {
+		t.Errorf("Publish() to empty channel = %d, want 0", count)
+	}
+}
+
+func TestHubChannelCount(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	if hub.ChannelCount("news") != 0 {
+		t.Error("new channel should have 0 subscribers")
+	}
+
+	hub.Subscribe(client, "news")
+	if hub.ChannelCount("news") != 1 {
+		t.Errorf("ChannelCount() = %d, want 1", hub.ChannelCount("news"))
+	}
+
+	hub.Subscribe(client, "news") // duplicate
+	if hub.ChannelCount("news") != 1 {
+		t.Errorf("ChannelCount() after duplicate = %d, want 1", hub.ChannelCount("news"))
+	}
+}
+
+func TestHubTotalChannels(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	if hub.TotalChannels() != 0 {
+		t.Errorf("TotalChannels() = %d, want 0", hub.TotalChannels())
+	}
+
+	hub.Subscribe(client, "news")
+	hub.Subscribe(client, "sports")
+
+	if hub.TotalChannels() != 2 {
+		t.Errorf("TotalChannels() = %d, want 2", hub.TotalChannels())
+	}
+}
+
+func TestHubChannelList(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	channels := hub.ChannelList("")
+	if len(channels) != 0 {
+		t.Errorf("Empty hub ChannelList() length = %d, want 0", len(channels))
+	}
+
+	hub.Subscribe(client, "news")
+	hub.Subscribe(client, "sports")
+
+	channels = hub.ChannelList("")
+	if len(channels) != 2 {
+		t.Errorf("ChannelList() length = %d, want 2", len(channels))
+	}
+}
+
+func TestHubMultipleClients(t *testing.T) {
+	hub := NewHub()
+	client1 := NewMockClient()
+	client2 := NewMockClient()
+	client3 := NewMockClient()
+
+	hub.Subscribe(client1, "news")
+	hub.Subscribe(client2, "news")
+	hub.Subscribe(client3, "news")
+
+	count := hub.ChannelCount("news")
+	if count != 3 {
+		t.Errorf("ChannelCount() = %d, want 3", count)
+	}
+}
+
+func TestHubChannelCleanup(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	hub.Subscribe(client, "temp")
+	if hub.ChannelCount("temp") != 1 {
+		t.Error("Channel should exist")
+	}
+
+	hub.Unsubscribe(client, "temp")
+	if hub.ChannelCount("temp") != 0 {
+		t.Error("Channel should be cleaned up")
+	}
+
+	channels := hub.ChannelList("")
+	for _, ch := range channels {
+		if ch == "temp" {
+			t.Error("Channel should not be in list after cleanup")
+		}
+	}
+}
+
+func TestNewClient(t *testing.T) {
+	conn := &mockConn{data: []byte{}}
+	client := NewClient(conn)
+	if client == nil {
+		t.Fatal("NewClient() returned nil")
+	}
+	if client.subscribed == nil {
+		t.Error("client.subscribed is nil")
+	}
+}
+
+func TestClientAddSubscription(t *testing.T) {
+	client := NewMockClient()
+
+	client.AddSubscription("news")
+	if client.SubsCount() != 1 {
+		t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
+	}
+
+	client.AddSubscription("sports")
+	if client.SubsCount() != 2 {
+		t.Errorf("SubsCount() = %d, want 2", client.SubsCount())
+	}
+}
+
+func TestClientRemoveSubscription(t *testing.T) {
+	client := NewMockClient()
+
+	client.AddSubscription("news")
+	client.AddSubscription("sports")
+	client.RemoveSubscription("news")
+
+	if client.SubsCount() != 1 {
+		t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
+	}
+}
+
+func TestClientGetChannels(t *testing.T) {
+	client := NewMockClient()
+
+	client.AddSubscription("news")
+	client.AddSubscription("sports")
+
+	channels := client.GetChannels()
+	if len(channels) != 2 {
+		t.Errorf("GetChannels() length = %d, want 2", len(channels))
+	}
+
+	found := make(map[string]bool)
+	for _, ch := range channels {
+		found[ch] = true
+	}
+	if !found["news"] || !found["sports"] {
+		t.Error("GetChannels() missing expected channels")
+	}
+}
+
+func TestClientGetChannelsReturnsCopy(t *testing.T) {
+	client := NewMockClient()
+	client.AddSubscription("news")
+
+	channels := client.GetChannels()
+	channels[0] = "modified"
+
+	channels2 := client.GetChannels()
+	if channels2[0] == "modified" {
+		t.Error("GetChannels() should return a copy, not modify original")
+	}
+}
+
+func TestClientSubsCount(t *testing.T) {
+	client := NewMockClient()
+
+	if client.SubsCount() != 0 {
+		t.Errorf("Initial SubsCount() = %d, want 0", client.SubsCount())
+	}
+
+	client.AddSubscription("news")
+	if client.SubsCount() != 1 {
+		t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
+	}
+}
+
+func TestConcurrentSubscribe(t *testing.T) {
+	hub := NewHub()
+	client := NewMockClient()
+
+	done := make(chan bool)
+	go func() {
+		hub.Subscribe(client, "channel1")
+		done <- true
+	}()
+
+	hub.Subscribe(client, "channel2")
+	<-done
+
+	if hub.TotalChannels() != 2 {
+		t.Errorf("TotalChannels() = %d, want 2", hub.TotalChannels())
+	}
+}