pekoe/internal/server/server.go

// @tag server
// @tag core
package server

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"math/rand"
	"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)

// generateRunID generates a 40-character hexadecimal run ID like Redis.
// Redis uses SHA1 of random data for 160 bits of randomness.
func generateRunID() string {
	b := make([]byte, 20)
	_, _ = rand.Read(b)
	return fmt.Sprintf("%x", b)
}

// Sentinel errors for connection handling
var (
	ErrQuit = errors.New("client quit")
)

// Version is set at build time via -ldflags
var Version = "unknown"

// 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
	version   string       // server version for INFO command
	startTime time.Time    // server start time for uptime calculation
	runID     string       // random run ID
}

// 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,
		version:   Version,
		startTime: time.Now(),
		runID:     generateRunID(),
	}
}

// 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())
	// Remove existing pid file first to handle stale files from previous runs
	// This also fixes permission denied errors when file already exists
	if err := os.Remove(s.cfg.PidFile); err != nil && !errors.Is(err, fs.ErrNotExist) {
		return fmt.Errorf("failed to remove existing pid file: %w", err)
	}
	//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)
	case "info":
		return s.handleInfo(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 with timeout
			shutdownTimeout := time.Duration(s.cfg.ShutdownTimeout) * time.Second
			if shutdownTimeout > 0 {
				// Create a channel to wait on wg.Wait() with timeout
				done := make(chan struct{})
				go func() {
					s.wg.Wait()
					close(done)
				}()

				select {
				case <-done:
					// All connections finished normally
				case <-time.After(shutdownTimeout):
					// Timeout reached - force cleanup
					logger.Warn("event=shutdown_timeout timeout=%ds active_connections=%d", s.cfg.ShutdownTimeout, s.active.Load())
				}
			} else {
				// No timeout - wait forever (original behavior)
				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
	}
}