// @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() }