// @tag server // @tag core package server import ( "context" "errors" "fmt" "io" "io/fs" "net" "os" "strconv" "strings" "sync" "sync/atomic" "time" "github.com/eleksir/pekoe/internal/config" "github.com/eleksir/pekoe/internal/logger" "github.com/eleksir/pekoe/internal/pubsub" ) const MaxDatabases = 16 // Redis default maximum databases (0-15) // Sentinel errors for connection handling var ( ErrQuit = errors.New("client quit") ) // Server is the main Pekoe server instance. // It manages the TCP listener, connection lifecycle, and delegates to the pubsub hub. type Server struct { cfg *config.Config hub pubsubHub ln net.Listener wg sync.WaitGroup ctx context.Context cancel context.CancelFunc active atomic.Int64 // atomic counter for active connections } // PubSubHub defines the interface our server expects from the pubsub implementation. // This allows for cleaner dependency injection and testing. type pubsubHub interface { Subscribe(client *pubsub.Client, channel string) bool Unsubscribe(client *pubsub.Client, channel string) bool UnsubscribeAll(client *pubsub.Client) Publish(channel string, message []byte) int PublishToPattern(channel string, message []byte) int ChannelCount(channel string) int ChannelList(pattern string) []string TotalChannels() int PatternSubscribe(client *pubsub.Client, pattern string) (bool, error) PatternUnsubscribe(client *pubsub.Client, pattern string) bool PatternCount() int RegisterClient(client *pubsub.Client) UnregisterClient(client *pubsub.Client) GetClient(id uint64) *pubsub.Client GetAllClients() []*pubsub.Client ClientCount() int } // New creates a new Pekoe server with the given config and pubsub hub. func New(cfg *config.Config, hub pubsubHub) *Server { ctx, cancel := context.WithCancel(context.Background()) return &Server{ cfg: cfg, hub: hub, ctx: ctx, cancel: cancel, } } // Run starts the server and blocks until shutdown. // Listens on the configured address, accepts connections, and handles them. // Returns an error only if the initial listen fails. func (s *Server) Run() error { ln, err := net.Listen("tcp", s.cfg.Addr()) if err != nil { return fmt.Errorf("failed to listen on %s: %w", s.cfg.Addr(), err) } s.ln = ln logger.Info("event=startup addr=%s max_clients=%d", s.cfg.Addr(), s.cfg.MaxClients) // Create PID file if configured if s.cfg.PidFile != "" { if err := s.createPidFile(); err != nil { logger.Error("event=pid_file_error error=%v", err) // Don't fail startup, just log the error } else { logger.Info("event=pid_file_created path=%s pid=%d", s.cfg.PidFile, os.Getpid()) } } s.wg.Add(1) go s.acceptLoop() // Wait for shutdown signal s.waitForSignal() return nil } // createPidFile creates the PID file with the current process ID func (s *Server) createPidFile() error { pid := strconv.Itoa(os.Getpid()) //nolint: gosec return os.WriteFile(s.cfg.PidFile, []byte(pid), 0644) } // removePidFile removes the PID file on shutdown func (s *Server) removePidFile() { if s.cfg.PidFile == "" { return } if err := os.Remove(s.cfg.PidFile); err != nil { if !errors.Is(err, fs.ErrNotExist) { logger.Error("event=pid_file_remove_error path=%s error=%v", s.cfg.PidFile, err) } } else { logger.Info("event=pid_file_removed path=%s", s.cfg.PidFile) } } // AcceptLoop runs in a goroutine and accepts incoming connections. // Uses a loop with error handling to continuously accept connections until shutdown. func (s *Server) acceptLoop() { defer s.wg.Done() for { conn, err := s.ln.Accept() if err != nil { // Check if we're shutting down (context cancelled) vs real error select { case <-s.ctx.Done(): return default: } // Handle temporary errors gracefully - they're often transient //nolint: staticcheck if ne, ok := err.(net.Error); ok && ne.Temporary() { logger.Warn("event=accept_temporary error=%v retry_ms=5", err) time.Sleep(5 * time.Millisecond) continue } logger.Error("event=accept_error error=%v", err) return } // Enforce max clients limit - protects server from resource exhaustion if s.active.Load() >= int64(s.cfg.MaxClients) { logger.Warn("event=max_clients_reached client=%s", conn.RemoteAddr()) _ = conn.Close() continue } go s.handleConn(conn) } } // HandleConn manages the lifecycle of a single client connection. // Parses commands, processes them, and sends responses. // Uses per-connection goroutine to allow concurrent client handling. // nolint: cyclop func (s *Server) handleConn(conn net.Conn) { s.wg.Add(1) defer s.wg.Done() s.active.Add(1) defer s.active.Add(-1) // Create and register client client := pubsub.NewClient(conn) s.hub.RegisterClient(client) defer func() { s.hub.UnregisterClient(client) s.hub.UnsubscribeAll(client) _ = client.Close() }() logger.Info("client=%s event=connected client_id=%d", client.RemoteAddr(), client.ID()) // Set timeouts if configured if s.cfg.ReadTimeoutDuration() > 0 { _ = conn.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeoutDuration())) } if s.cfg.WriteTimeoutDuration() > 0 { _ = conn.SetWriteDeadline(time.Now().Add(s.cfg.WriteTimeoutDuration())) } reader := NewConnReader(conn, s.cfg.MaxMsgBytes) writer := NewConnWriter(conn) // Main command loop - process commands until error or disconnect for { // Check for shutdown before each command select { case <-s.ctx.Done(): return default: } // Reset read deadline before each command if using timeouts if s.cfg.ReadTimeoutDuration() > 0 { _ = conn.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeoutDuration())) } cmd, err := reader.ReadCommand() if err != nil { if err == io.EOF { logger.Info("client=%s event=disconnected client_id=%d", client.RemoteAddr(), client.ID()) return // Client disconnected normally } logger.Error("client=%s error=protocol message=%v", client.RemoteAddr(), err) writer.WriteError("Protocol error: " + err.Error()) _ = writer.Flush() return } // Log incoming command at DEBUG level cmdName := "" if len(cmd.Args) > 0 { cmdName = strings.ToLower(string(cmd.Args[0])) } logger.Debug("client=%s command=%s args=%d", client.RemoteAddr(), cmdName, len(cmd.Args)-1) // Process command - handler returns error only on fatal issues if err := s.processCommand(client, cmd, writer); err != nil { if errors.Is(err, ErrQuit) { _ = writer.Flush() logger.Info("client=%s event=quit", client.RemoteAddr()) return // Graceful client disconnect via QUIT } logger.Error("client=%s error=processing message=%v", client.RemoteAddr(), err) return } // Flush response - important for pubsub as we need immediate feedback if s.cfg.WriteTimeoutDuration() > 0 { _ = conn.SetWriteDeadline(time.Now().Add(s.cfg.WriteTimeoutDuration())) } if err := writer.Flush(); err != nil { logger.Error("client=%s error=flush message=%v", client.RemoteAddr(), err) return } } } // processCommand routes commands to their handlers. // Each command type has its own handler for clarity. // nolint: cyclop func (s *Server) processCommand(client *pubsub.Client, cmd *Command, writer *ConnWriter) error { if cmd == nil || len(cmd.Args) == 0 { writer.WriteError("ERR empty command") return nil } // Convert command name to lowercase for case-insensitive matching (Redis standard) name := strings.ToLower(string(cmd.Args[0])) switch name { case "subscribe": return s.handleSubscribe(client, cmd.Args[1:], writer) case "unsubscribe": return s.handleUnsubscribe(client, cmd.Args[1:], writer) case "psubscribe": return s.handlePsubscribe(client, cmd.Args[1:], writer) case "punsubscribe": return s.handlePunsubscribe(client, cmd.Args[1:], writer) case "publish": return s.handlePublish(cmd.Args[1:], writer) case "pubsub": return s.handlePubsub(cmd.Args[1:], writer) case "ping": return s.handlePing(cmd.Args[1:], writer) case "quit": return s.handleQuit(client, writer) case "client": return s.handleClient(client, cmd.Args[1:], writer) case "select": return s.handleSelect(client, cmd.Args[1:], writer) default: // Check if client is in pubsub mode - only pubsub commands allowed if client != nil && (client.SubsCount() > 0 || client.PatternSubsCount() > 0) { logger.Error("client=%s error=invalid_in_pubsub_mode command=%s", client.RemoteAddr(), name) writer.WriteError("ERR Can't execute commands in pubsub mode") return nil } clientAddr := "" if client != nil { clientAddr = client.RemoteAddr() } logger.Error("client=%s error=unknown_command command=%s", clientAddr, name) writer.WriteError(fmt.Sprintf("ERR unknown command '%s'", name)) return nil } } // waitForSignal listens for signals and handles them appropriately. // On Unix: handles SIGINT, SIGTERM (shutdown), SIGHUP (reserved), SIGUSR1 (log reopen) // On Windows: handles SIGINT, SIGTERM (shutdown), SIGHUP (log reopen) func (s *Server) waitForSignal() { sigCh := signalSetup() for { sig := <-sigCh if s.handleSignal(sig) { // Shutdown signal received - perform graceful shutdown s.cancel() // Close listener - causes Accept() to return error, stopping accept loop if err := s.ln.Close(); err != nil { logger.Error("event=listener_close_error error=%v", err) } // Wait for all active connections to finish s.wg.Wait() // Remove PID file s.removePidFile() logger.Info("event=shutdown_complete active_connections=%d", s.active.Load()) return } // Not a shutdown signal - continue listening } }