Initial commit
This commit is contained in:
eleksir
325
internal/server/server.go
Normal file
325
internal/server/server.go
Normal file
@@ -0,0 +1,325 @@
|
||||
// @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"
|
||||
)
|
||||
|
||||
// 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)
|
||||
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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user