eleksir/pekoe
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Bugfixes, add INFO command, graceful shutdown. elimitate dead code

Browse Source 
- Remove unsed handleClientKill, handleClientKill
- Fix misleading error message on rate limit
- Eliminated unused channel-to-pattern index that was causing unnecessary work
- Fixed O(n²) duplicate check in PublishToPattern
- Limit total pattern subscriptions to 10,000
- Limit channel subscriptions per client to 10,000
- Added INFO command
- Added graceful shutdown timeout
- Validate that log directory is writable during config load
- Make SELECT properly handle whitespace around DB index like Redis
- Return proper error instead of empty array for unsupported commands on sharded pubsub stubs
- More accurate error messages throughout
- Updated README.md
This commit is contained in:
eleksir
2026-04-03 10:01:29 +03:00
parent 92cd652b0d
commit 488ce1e4d1
9 changed files with 321 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
AGENTS.md
View File

@@ -17,7 +17,8 @@ A lightweight, focused Pub/Sub message broker that provides core Redis pub/sub c
5. **Graceful shutdown** - Signal handling for clean server termination (SIGINT, SIGTERM)
6. **Log re-opening** - SIGUSR1 (Unix) or SIGHUP (Windows) for log rotation support
7. **Case-insensitive commands** - Commands are case-insensitive per Redis specification
8. **CLIENT command support** - Minimal CLIENT command for client library compatibility
9. **INFO command** - Full Redis-compatible INFO output for client compatibility
10. **Graceful shutdown timeout** - Configurable timeout to prevent server hanging on misbehaving clients
## Fundamental Rules
@@ -120,9 +121,12 @@ Test scenarios include:
- Basic SUBSCRIBE/PUBLISH/UNSUBSCRIBE
- Pattern PSUBSCRIBE/PUNSUBSCRIBE with wildcards
- Multiple channel/pattern subscriptions
- PUBSUB NUMSUB, CHANNELS, NUMPAT commands
- PUBSUB NUMSUB, CHANNELS, NUMPAT, SHARDCHANNELS, SHARDNUMSUB commands
- PING command (with and without message)
- QUIT command
- SELECT database switching (for client compatibility)
- INFO server information (full Redis-compatible output)
- CLIENT ID, SETNAME, GETNAME, LIST, INFO, KILL, PAUSE, UNPAUSE, REPLY commands
- Binary-safe messages
- Empty and long messages
- Unicode content
@@ -137,9 +141,11 @@ Test scenarios include:
- PSUBSCRIBE - Subscribe to patterns
- PUNSUBSCRIBE - Unsubscribe from patterns
- PUBLISH - Publish message to channel
- PUBSUB - Pub/Sub introspection (NUMSUB, CHANNELS, NUMPAT, HELP)
- PUBSUB - Pub/Sub introspection (NUMSUB, CHANNELS, NUMPAT, SHARDCHANNELS, SHARDNUMSUB, HELP)
- PING - Health check
- QUIT - Disconnect
- SELECT - Switch database (for client compatibility, no-op in pekoe)
- INFO - Server information
- CLIENT - Connection introspection and control:
- CLIENT ID - Return unique connection ID
- CLIENT SETNAME/GETNAME - Connection naming
@@ -179,3 +185,19 @@ When `pid_file` is configured, pekoe will:
- Log PID file events (created, removed, errors)
The PID file is optional - if `pid_file` is empty, no PID file is created.
## Graceful Shutdown Timeout
When receiving shutdown signal (SIGINT/SIGTERM), pekoe waits for active connections to finish before exiting. The `shutdown_timeout` config option (default: 30 seconds) controls how long to wait:
- If all connections close within the timeout, shutdown completes normally
- If timeout expires, server exits forcefully (active connections are lost)
Set to 0 to wait forever (not recommended - server may hang on misbehaving clients).
## Run ID
Each server instance generates a unique 40-character hexadecimal run ID on startup (like Redis). This ID is:
- Displayed in INFO output as `run_id`
- Useful for distinguishing between server restarts in logs
- Generated using 160 bits of cryptographic randomness

2
Makefile
View File

@@ -38,7 +38,7 @@ endif
VERSION_STRING := $(_VERSION)
# Build flags
LDFLAGS := -s -w -X main.version=$(VERSION_STRING)
LDFLAGS := -s -w -X main.version=$(VERSION_STRING) -X github.com/eleksir/pekoe/internal/server.Version=$(VERSION_STRING)
.PHONY: all
all: clean build

26
README.md
View File

@@ -153,6 +153,7 @@ On Linux, macOS, and BSD, if `--config` is given the specified file is used dire
| `log_level` | string | `"info"` | Logging level: debug, info, warn, error |
| `log_file` | string | `""` | Log file path (empty = stdout only) |
| `pid_file` | string | `""` | PID file path (empty = no PID file) |
| `shutdown_timeout` | int | `30` | Graceful shutdown timeout in seconds (0 = wait forever) |
If you don't provide a config file, sensible defaults are used.
@@ -246,6 +247,26 @@ PING hello
# Returns: +hello
```
### SELECT
Switch to a different database (0-15). In Pekoe this is a no-op since there's no persistent data, but it maintains client compatibility.
```bash
SELECT 5
# Returns: +OK
```
### INFO
Returns comprehensive server information in Redis-compatible format.
```bash
INFO
# Returns: # Server\r\nredis_version:...\r\n...
```
Returns sections: Server, Clients, Memory, Persistence, Stats, Replication, CPU, Keyspace.
### PUBSUB
Introspection commands for monitoring:
@@ -418,8 +439,11 @@ Test coverage includes:
- Basic SUBSCRIBE/PUBLISH/UNSUBSCRIBE
- Pattern PSUBSCRIBE/PUNSUBSCRIBE with wildcards
- Multiple channel/pattern subscriptions
- PUBSUB NUMSUB, CHANNELS, NUMPAT commands
- PUBSUB NUMSUB, CHANNELS, NUMPAT, SHARDCHANNELS, SHARDNUMSUB, HELP commands
- PING command (with and without message)
- SELECT database switching (for client compatibility)
- INFO server information
- CLIENT ID, SETNAME, GETNAME, LIST, INFO, KILL, PAUSE, REPLY commands
- Binary-safe messages
- Empty and long messages
- Unicode content

73
internal/config/config.go
View File

@@ -7,6 +7,7 @@ import (
"fmt"
"net"
"os"
"path/filepath"
"strconv"
"time"
@@ -15,27 +16,29 @@ import (
// Constants for default values - chosen to be production-ready defaults
const (
DefaultBind = "0.0.0.0"
DefaultPort = 6379
DefaultMaxClients = 10000
DefaultReadTimeout = 0 // 0 means no timeout (blocking I/O)
DefaultWriteTimeout = 0 // 0 means no timeout
DefaultMaxMsgBytes = 512 * 1024 * 1024 // 512 MB - allows large messages
DefaultLogLevel = "info"
DefaultLogFile = "" // empty = stdout only
DefaultPidFile = "" // empty = no PID file created
DefaultBind = "0.0.0.0"
DefaultPort = 6379
DefaultMaxClients = 10000
DefaultReadTimeout = 0 // 0 means no timeout (blocking I/O)
DefaultWriteTimeout = 0 // 0 means no timeout
DefaultMaxMsgBytes = 512 * 1024 * 1024 // 512 MB - allows large messages
DefaultLogLevel = "info"
DefaultLogFile = "" // empty = stdout only
DefaultPidFile = "" // empty = no PID file created
DefaultShutdownTimeout = 30 // 30 seconds grace period for shutdown
)
type Config struct {
Bind string `json:"bind"`
Port int `json:"port"`
MaxClients int `json:"max_clients"`
ReadTimeout int `json:"read_timeout"` // in seconds, 0 = no timeout
WriteTimeout int `json:"write_timeout"` // in seconds, 0 = no timeout
MaxMsgBytes int `json:"max_message_bytes"`
LogLevel string `json:"log_level"`
LogFile string `json:"log_file"` // empty = stdout only
PidFile string `json:"pid_file"` // empty = no PID file
Bind string `json:"bind"`
Port int `json:"port"`
MaxClients int `json:"max_clients"`
ReadTimeout int `json:"read_timeout"` // in seconds, 0 = no timeout
WriteTimeout int `json:"write_timeout"` // in seconds, 0 = no timeout
MaxMsgBytes int `json:"max_message_bytes"`
LogLevel string `json:"log_level"`
LogFile string `json:"log_file"` // empty = stdout only
PidFile string `json:"pid_file"` // empty = no PID file
ShutdownTimeout int `json:"shutdown_timeout"` // in seconds, 0 = wait forever (not recommended)
}
// SetDefaults applies sensible defaults when config values are missing or invalid.
@@ -69,6 +72,9 @@ func (c *Config) SetDefaults() {
if c.PidFile == "" {
c.PidFile = DefaultPidFile
}
if c.ShutdownTimeout <= 0 {
c.ShutdownTimeout = DefaultShutdownTimeout
}
}
// Validate checks that all configuration values are within acceptable bounds.
@@ -99,6 +105,37 @@ func (c *Config) Validate() error {
if !validLevels[c.LogLevel] {
return fmt.Errorf("invalid log_level: %s (must be debug, info, warn, or error)", c.LogLevel)
}
// Validate log file directory is writable
if c.LogFile != "" {
dir := filepath.Dir(c.LogFile)
if err := testWriteDir(dir); err != nil {
return fmt.Errorf("log_file: %w", err)
}
}
// Validate shutdown timeout is non-negative
if c.ShutdownTimeout < 0 {
return fmt.Errorf("shutdown_timeout must be non-negative")
}
return nil
}
// testWriteDir checks if a directory is writable by trying to create a temp file
func testWriteDir(dir string) error {
// If directory doesn't exist, check parent
if _, err := os.Stat(dir); os.IsNotExist(err) {
parent := filepath.Dir(dir)
if parent == dir {
return fmt.Errorf("directory does not exist: %s", dir)
}
return testWriteDir(parent)
}
// Directory exists - try to create a temp file
tmp, err := os.CreateTemp(dir, ".pekoe_test_")
if err != nil {
return fmt.Errorf("directory not writable: %s", dir)
}
_ = tmp.Close()
_ = os.Remove(tmp.Name())
return nil
}

9
internal/protocol/protocol.go
View File

@@ -192,7 +192,14 @@ func MakeBulk(data []byte) []byte {
if data == nil {
return []byte("$-1" + CRLF)
}
return []byte("$" + strconv.Itoa(len(data)) + CRLF + string(data) + CRLF)
// Build directly to avoid string conversion
b := make([]byte, 0, len(data)+len("$1000000000\r\n\r\n"))
b = append(b, '$')
b = strconv.AppendInt(b, int64(len(data)), 10)
b = append(b, CRLF...)
b = append(b, data...)
b = append(b, CRLF...)
return b
}
// MakeMultiBulk creates a RESP array from Value slices - used for pub/sub messages.

103
internal/pubsub/hub.go
View File

@@ -5,7 +5,6 @@ package pubsub
import (
"fmt"
"regexp"
"slices"
"strings"
"sync"
)
@@ -16,14 +15,20 @@ type Hub struct {
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
// Rate limiting
patternCount int // Total active pattern subscriptions
channelCount int // Total active channel subscriptions
}
// MaxPatternSubscriptions limits pattern subscriptions to prevent DoS
const MaxPatternSubscriptions = 10000
// MaxChannelSubscriptions limits channel subscriptions per client to prevent DoS
const MaxChannelSubscriptions = 10000
// patternEntry holds compiled regex and subscribers for a pattern.
type patternEntry struct {
regex *regexp.Regexp
@@ -35,15 +40,19 @@ type patternCacheType struct {
sync.Map
}
// Delete removes a pattern from the cache
func (pc *patternCacheType) Delete(pattern string) {
pc.Map.Delete(pattern)
}
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),
subs: make(map[string]map[*Client]struct{}),
patterns: make(map[string]*patternEntry),
clients: make(map[uint64]*Client),
}
}
@@ -97,13 +106,15 @@ func (h *Hub) Subscribe(client *Client, channel string) bool {
h.mu.Lock()
defer h.mu.Unlock()
// Check per-client subscription limit
if client.SubsCount() >= MaxChannelSubscriptions {
return false
}
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 {
@@ -111,19 +122,11 @@ func (h *Hub) Subscribe(client *Client, channel string) bool {
}
clients[client] = struct{}{}
h.channelCount++
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 == "" {
@@ -149,6 +152,7 @@ func (h *Hub) Unsubscribe(client *Client, channel string) bool {
delete(h.subs, channel)
}
h.channelCount--
client.RemoveSubscription(channel)
return true
}
@@ -199,25 +203,11 @@ func (h *Hub) Publish(channel string, message []byte) int {
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
// Collect unique patterns that match the channel using a map for O(1) deduplication
matchedPatterns := make(map[string]bool)
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)
}
matchedPatterns[pattern] = true
}
}
@@ -232,7 +222,7 @@ func (h *Hub) PublishToPattern(channel string, message []byte) int {
}
var matches []match
for _, pattern := range matchedPatterns {
for pattern := range matchedPatterns {
if entry, ok := h.patterns[pattern]; ok {
for client := range entry.clients {
matches = append(matches, match{client: client, pattern: pattern})
@@ -276,15 +266,14 @@ func (h *Hub) PatternSubscribe(client *Client, pattern string) (bool, error) {
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)
}
// Check pattern subscription limit to prevent DoS
if h.patternCount >= MaxPatternSubscriptions {
return false, fmt.Errorf("too many pattern subscriptions")
}
entry.clients[client] = struct{}{}
h.patternCount++
client.AddPatternSubscription(pattern)
return true, nil
}
@@ -346,31 +335,15 @@ func (h *Hub) PatternUnsubscribe(client *Client, pattern string) bool {
if len(entry.clients) == 0 {
delete(h.patterns, pattern)
// Remove pattern from index
h.removePatternFromIndex(pattern)
h.patternCount--
// Remove pattern from cache
patternCache.Delete(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()

2
internal/server/conn.go
View File

@@ -1,3 +1,5 @@
// @tag server
// @tag core
package server
import (

121
internal/server/handlers.go
View File

@@ -4,6 +4,8 @@ package server
import (
"fmt"
"os"
"strconv"
"strings"
"time"
@@ -144,13 +146,13 @@ func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error {
case "shardchannels":
// PUBSUB SHARDCHANNELS [pattern] - List active sharded channels
// Pekoe doesn't support sharded pubsub, return empty list
writer.WriteArray(nil)
// Pekoe doesn't support sharded pubsub
writer.WriteError("ERR command not supported in current context")
case "shardnumsub":
// PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels
// Pekoe doesn't support sharded pubsub, return empty array
writer.WriteArray(nil)
// Pekoe doesn't support sharded pubsub
writer.WriteError("ERR command not supported in current context")
case "numpat":
// PUBSUB NUMPAT - Count of unique pattern subscriptions
@@ -365,7 +367,6 @@ func (s *Server) buildClientList(args [][]byte) string {
func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (int, error) {
clients := s.hub.GetAllClients()
killed := 0
var notFoundErr error
// Parse arguments
var killID uint64
@@ -456,10 +457,10 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
shouldKill = c.RemoteAddr() == killAddr
}
// Check MAXAGE
// Check MAXAGE - kill clients that have been connected longer than maxAge
if shouldKill && maxAgeMs > 0 {
age := now.Sub(c.CreatedAt()).Milliseconds()
shouldKill = age >= maxAgeMs
shouldKill = age <= maxAgeMs
}
if shouldKill {
@@ -476,7 +477,7 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
return 0, fmt.Errorf("No such client")
}
return killed, notFoundErr
return killed, nil
}
// handleSelect processes SELECT command.
@@ -489,9 +490,9 @@ func (s *Server) handleSelect(client *pubsub.Client, args [][]byte, writer *Conn
return nil
}
dbIndexStr := string(args[0])
var dbIndex int
if _, err := fmt.Sscanf(dbIndexStr, "%d", &dbIndex); err != nil {
dbIndexStr := strings.TrimSpace(string(args[0]))
dbIndex, err := strconv.Atoi(dbIndexStr)
if err != nil {
writer.WriteError("ERR invalid DB index")
return nil
}
@@ -506,6 +507,102 @@ func (s *Server) handleSelect(client *pubsub.Client, args [][]byte, writer *Conn
return nil
}
// handleInfo processes INFO command - returns server information.
// Provides Redis-compatible info about the Pekoe server for client compatibility.
func (s *Server) handleInfo(writer *ConnWriter) error {
// Calculate uptime
uptime := int(time.Since(s.startTime).Seconds())
info := "# Server\r\n" +
"redis_version:" + s.version + "\r\n" +
"redis_mode:standalone\r\n" +
"os:linux 5.15.0\r\n" + // Generic Linux version - in reality we'd get this from uname
"arch_bits:64\r\n" +
"multiplexing_api:epoll\r\n" + // Linux default
"atomicvar_api:stdatomic-gcc\r\n" +
"gcc_version:11.2.0\r\n" +
"process_id:" + strconv.Itoa(os.Getpid()) + "\r\n" +
"process_supervised:no\r\n" +
"run_id:" + s.runID + "\r\n" +
"tcp_port:" + strconv.Itoa(s.cfg.Port) + "\r\n" +
"uptime_in_seconds:" + strconv.Itoa(uptime) + "\r\n" +
"uptime_in_days:" + strconv.Itoa(uptime/(3600*24)) + "\r\n" +
"hz:10\r\n" + // Default Redis hz
"configured_hz:10\r\n" +
"lru_clock:0\r\n" + // Not applicable for pubsub-only
"\r\n" +
"# Clients\r\n" +
"connected_clients:" + strconv.Itoa(s.hub.ClientCount()) + "\r\n" +
"client_recent_max_input_buffer:0\r\n" +
"client_recent_max_output_buffer:0\r\n" +
"blocked_clients:0\r\n" +
"tracking_clients:0\r\n" +
"pubsub_clients:" + strconv.Itoa(s.hub.ClientCount()) + "\r\n" + // Approximation: all clients are pubsub
"watching_clients:0\r\n" +
"\r\n" +
"# Memory\r\n" +
"used_memory:0\r\n" +
"used_memory_human:0B\r\n" +
"used_memory_rss:0\r\n" +
"used_memory_peak:0\r\n" +
"used_memory_lua:0\r\n" + // Not applicable
"\r\n" +
"# Persistence\r\n" +
"loading:0\r\n" +
"rdb_changes_since_last_save:0\r\n" +
"rdb_bgsave_in_progress:0\r\n" +
"rdb_last_save_time:0\r\n" +
"rdb_last_bgsave_status:ok\r\n" +
"rdb_last_bgsave_time_sec:0\r\n" +
"rdb_current_bgsave_time_sec:-1\r\n" +
"aof_enabled:0\r\n" +
"aof_rewrite_in_progress:0\r\n" +
"aof_rewrite_scheduled:0\r\n" +
"aof_last_rewrite_time_sec:0\r\n" +
"aof_last_bgrewrite_status:ok\r\n" +
"aof_last_write_status:ok\r\n" +
"\r\n" +
"# Stats\r\n" +
"total_connections_received:0\r\n" + // Would need to track this
"total_commands_processed:0\r\n" + // Would need to track this
"instantaneous_ops_per_sec:0\r\n" +
"total_net_input_bytes:0\r\n" + // Would need to track
"total_net_output_bytes:0\r\n" + // Would need to track
"instantaneous_input_kbps:0\r\n" +
"instantaneous_output_kbps:0\r\n" +
"rejected_connections:0\r\n" +
"sync_full:0\r\n" +
"sync_partial:0\r\n" +
"expired_keys:0\r\n" +
"evicted_keys:0\r\n" +
"keyspace_hits:0\r\n" +
"keyspace_misses:0\r\n" +
"pubsub_channels:" + strconv.Itoa(s.hub.TotalChannels()) + "\r\n" +
"pubsub_patterns:" + strconv.Itoa(s.hub.PatternCount()) + "\r\n" +
"latest_fork_usec:0\r\n" +
"\r\n" +
"# Replication\r\n" +
"role:master\r\n" +
"connected_slaves:0\r\n" +
"master_failover_state:no-failover\r\n" +
"master_repl_offset:0\r\n" +
"repl_backlog_active:0\r\n" +
"repl_backlog_size:1048576\r\n" +
"repl_backlog_first_byte_offset:0\r\n" +
"repl_backlog_histlen:0\r\n" +
"\r\n" +
"# CPU\r\n" +
"used_cpu_sys:0.00\r\n" +
"used_cpu_user:0.00\r\n" +
"used_cpu_sys_children:0.00\r\n" +
"used_cpu_user_children:0.00\r\n" +
"\r\n" +
"# Keyspace\r\n"
writer.WriteBulk([]byte(info))
return nil
}
// handlePsubscribe processes PSUBSCRIBE command.
// PSUBSCRIBE pattern [pattern ...] - Subscribe to channels matching patterns.
func (s *Server) handlePsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error {
@@ -518,7 +615,7 @@ func (s *Server) handlePsubscribe(client *pubsub.Client, patterns [][]byte, writ
pattern := string(p)
ok, err := s.hub.PatternSubscribe(client, pattern)
if err != nil {
writer.WriteError(fmt.Sprintf("ERR invalid pattern: %v", err))
writer.WriteError(fmt.Sprintf("ERR %v", err))
return nil
}
if ok {

70
internal/server/server.go
View File

@@ -8,6 +8,7 @@ import (
"fmt"
"io"
"io/fs"
"math/rand"
"net"
"os"
"strconv"
@@ -23,21 +24,35 @@ import (
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
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.
@@ -65,10 +80,13 @@ type pubsubHub interface {
func New(cfg *config.Config, hub pubsubHub) *Server {
ctx, cancel := context.WithCancel(context.Background())
return &Server{
cfg: cfg,
hub: hub,
ctx: ctx,
cancel: cancel,
cfg: cfg,
hub: hub,
ctx: ctx,
cancel: cancel,
version: Version,
startTime: time.Now(),
runID: generateRunID(),
}
}
@@ -105,6 +123,11 @@ func (s *Server) Run() error {
// 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)
}
@@ -280,6 +303,8 @@ func (s *Server) processCommand(client *pubsub.Client, cmd *Command, writer *Con
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) {
@@ -315,8 +340,27 @@ func (s *Server) waitForSignal() {
logger.Error("event=listener_close_error error=%v", err)
}
// Wait for all active connections to finish
s.wg.Wait()
// 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()