Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
488ce1e4d1 | ||
|
|
92cd652b0d |
28
AGENTS.md
28
AGENTS.md
@@ -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
2
Makefile
@@ -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
26
README.md
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
58
internal/glob/glob.go
Normal file
58
internal/glob/glob.go
Normal file
@@ -0,0 +1,58 @@
|
||||
// @tag glob
|
||||
package glob
|
||||
|
||||
// Match checks if a string matches a glob-style pattern.
|
||||
// Supports:
|
||||
// - * (asterisk): matches any sequence of characters (including empty)
|
||||
// - ? (question): matches any single character
|
||||
//
|
||||
// This matches Redis PUBSUB CHANNELS pattern behavior.
|
||||
func Match(pattern, s string) bool {
|
||||
if pattern == "" {
|
||||
return s == ""
|
||||
}
|
||||
if pattern == "*" {
|
||||
return true
|
||||
}
|
||||
|
||||
// Simple glob matching using dynamic programming
|
||||
// Convert pattern to regex-like matching
|
||||
patternLen := len(pattern)
|
||||
sLen := len(s)
|
||||
|
||||
// dp[i][j] = true if pattern[0:i] matches s[0:j]
|
||||
dp := make([][]bool, patternLen+1)
|
||||
for i := range dp {
|
||||
dp[i] = make([]bool, sLen+1)
|
||||
}
|
||||
|
||||
// Empty pattern matches empty string
|
||||
dp[0][0] = true
|
||||
|
||||
// Handle leading * in pattern
|
||||
for i := 1; i <= patternLen; i++ {
|
||||
if pattern[i-1] == '*' {
|
||||
dp[i][0] = dp[i-1][0]
|
||||
}
|
||||
}
|
||||
|
||||
// Fill the DP table
|
||||
for i := 1; i <= patternLen; i++ {
|
||||
for j := 1; j <= sLen; j++ {
|
||||
switch pattern[i-1] {
|
||||
case '*':
|
||||
// * matches: either consume nothing (stay at same position in pattern)
|
||||
// or consume one character from string
|
||||
dp[i][j] = dp[i-1][j] || dp[i][j-1]
|
||||
case '?':
|
||||
// ? matches any single character
|
||||
dp[i][j] = dp[i-1][j-1]
|
||||
default:
|
||||
// Literal character must match exactly
|
||||
dp[i][j] = dp[i-1][j-1] && pattern[i-1] == s[j-1]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return dp[patternLen][sLen]
|
||||
}
|
||||
69
internal/glob/glob_test.go
Normal file
69
internal/glob/glob_test.go
Normal file
@@ -0,0 +1,69 @@
|
||||
// @tag glob
|
||||
// @tag test
|
||||
package glob
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestMatch(t *testing.T) {
|
||||
tests := []struct {
|
||||
pattern string
|
||||
s string
|
||||
want bool
|
||||
}{
|
||||
{"", "", true},
|
||||
{"", "foo", false},
|
||||
{"*", "", true},
|
||||
{"*", "foo", true},
|
||||
{"*", "foo/bar", true},
|
||||
{"foo", "foo", true},
|
||||
{"foo", "bar", false},
|
||||
{"foo*", "foo", true},
|
||||
{"foo*", "foobar", true},
|
||||
{"foo*", "fo", false},
|
||||
{"*foo", "foo", true},
|
||||
{"*foo", "barfoo", true},
|
||||
{"*foo", "bar", false},
|
||||
{"foo*bar", "foobar", true},
|
||||
{"foo*bar", "foobazbar", true},
|
||||
{"foo*bar", "foobaz", false},
|
||||
{"foo?bar", "fooXbar", true},
|
||||
{"foo?bar", "fooXXbar", false},
|
||||
{"foo??bar", "fooXXbar", true},
|
||||
{"?", "a", true},
|
||||
{"?", "", false},
|
||||
{"a?c", "abc", true},
|
||||
{"a?c", "ac", false},
|
||||
{"foo?*bar", "fooabcbar", true},
|
||||
{"foo?*bar", "foobar", false},
|
||||
{"*foo*bar*", "ahellofoocherrybarb", true},
|
||||
{"ch*", "channel", true},
|
||||
{"ch*", "ch", true},
|
||||
{"ch*", "c", false},
|
||||
{"*", "*", true},
|
||||
{"**", "*", true},
|
||||
{"*a*b", "aXb", true},
|
||||
{"a*b*c", "aXbYc", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.pattern+"/"+tt.s, func(t *testing.T) {
|
||||
if got := Match(tt.pattern, tt.s); got != tt.want {
|
||||
t.Errorf("Match(%q, %q) = %v, want %v", tt.pattern, tt.s, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkMatch(b *testing.B) {
|
||||
patterns := []string{"*", "ch*", "foo*bar", "???", "*channel*"}
|
||||
strings := []string{"channel", "channel123", "foobar", "abc", "somechannelname"}
|
||||
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
for _, p := range patterns {
|
||||
for _, s := range strings {
|
||||
Match(p, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -31,6 +31,7 @@ type Client struct {
|
||||
conn net.Conn
|
||||
name string // Connection name set via CLIENT SETNAME
|
||||
replyMode ReplyMode // Reply mode (ON/OFF/SKIP)
|
||||
selectedDB int // Selected database index (via SELECT command)
|
||||
flags uint64 // Client flags (e.g., paused, no-touch)
|
||||
created time.Time // When the client connected
|
||||
lastActive time.Time // Last activity time
|
||||
@@ -48,6 +49,7 @@ func NewClient(conn net.Conn) *Client {
|
||||
conn: conn,
|
||||
name: "",
|
||||
replyMode: ReplyOn,
|
||||
selectedDB: 0,
|
||||
flags: 0,
|
||||
created: now,
|
||||
lastActive: now,
|
||||
@@ -89,6 +91,27 @@ func (c *Client) GetReplyMode() ReplyMode {
|
||||
return c.replyMode
|
||||
}
|
||||
|
||||
// SelectedDB returns the currently selected database index
|
||||
func (c *Client) SelectedDB() int {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.selectedDB
|
||||
}
|
||||
|
||||
// CreatedAt returns when the client connected
|
||||
func (c *Client) CreatedAt() time.Time {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.created
|
||||
}
|
||||
|
||||
// SetSelectedDB sets the selected database index
|
||||
func (c *Client) SetSelectedDB(db int) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.selectedDB = db
|
||||
}
|
||||
|
||||
// UpdateActivity updates the last activity timestamp
|
||||
func (c *Client) UpdateActivity() {
|
||||
c.mu.Lock()
|
||||
@@ -212,7 +235,7 @@ func (c *Client) 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=
|
||||
// Format: id= addr= laddr= fd= age= idle= flags= db= sub= psub= multi= qbuf= qbuf-free= obl= oll= omem= tot-mem= events= cmd= user= redir= resp=
|
||||
func (c *Client) InfoString() string {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
@@ -229,16 +252,21 @@ func (c *Client) InfoString() string {
|
||||
flags += "S"
|
||||
}
|
||||
|
||||
// Events - indicate connection is active for read and write
|
||||
events := "rw"
|
||||
|
||||
// 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",
|
||||
return fmt.Sprintf("id=%d addr=%s laddr=%s fd=-1 age=%d idle=%d flags=%s db=%d sub=%d psub=%d multi=-1 qbuf=0 qbuf-free=0 obl=0 oll=0 omem=0 tot-mem=0 events=%s cmd=client user=default redir=-1 resp=2",
|
||||
c.id,
|
||||
c.RemoteAddr(),
|
||||
c.LocalAddr(),
|
||||
age,
|
||||
idle,
|
||||
flags,
|
||||
c.selectedDB,
|
||||
len(c.subscribed),
|
||||
len(c.patterns),
|
||||
events,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
// @tag server
|
||||
// @tag core
|
||||
package server
|
||||
|
||||
import (
|
||||
|
||||
@@ -4,8 +4,12 @@ package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/eleksir/pekoe/internal/glob"
|
||||
"github.com/eleksir/pekoe/internal/logger"
|
||||
"github.com/eleksir/pekoe/internal/protocol"
|
||||
"github.com/eleksir/pekoe/internal/pubsub"
|
||||
@@ -132,12 +136,24 @@ func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error {
|
||||
pattern = string(args[1])
|
||||
}
|
||||
channels := s.channelList(pattern)
|
||||
result := make([]protocol.Value, len(channels))
|
||||
for i, ch := range channels {
|
||||
result[i] = protocol.Bulk(ch)
|
||||
result := make([]protocol.Value, 0, len(channels))
|
||||
for _, ch := range channels {
|
||||
if pattern == "" || glob.Match(pattern, ch) {
|
||||
result = append(result, protocol.Bulk(ch))
|
||||
}
|
||||
}
|
||||
writer.WriteArray(result)
|
||||
|
||||
case "shardchannels":
|
||||
// PUBSUB SHARDCHANNELS [pattern] - List active sharded channels
|
||||
// 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
|
||||
writer.WriteError("ERR command not supported in current context")
|
||||
|
||||
case "numpat":
|
||||
// PUBSUB NUMPAT - Count of unique pattern subscriptions
|
||||
writer.WriteInt(int64(s.hub.PatternCount()))
|
||||
@@ -147,7 +163,9 @@ func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error {
|
||||
helpText := []string{
|
||||
"PUBSUB <subcommand> [<argument> [<argument> ...]]>",
|
||||
" PUBSUB CHANNELS [pattern] - List active channels (optionally filtered by pattern)",
|
||||
" PUBSUB SHARDCHANNELS [pattern] - List active sharded channels (not supported)",
|
||||
" PUBSUB NUMSUB [channel ...] - Get subscriber count for channels",
|
||||
" PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels (not supported)",
|
||||
" PUBSUB NUMPAT - Get the number of pattern subscriptions",
|
||||
" PUBSUB HELP - Show this help message",
|
||||
}
|
||||
@@ -352,8 +370,11 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
|
||||
|
||||
// Parse arguments
|
||||
var killID uint64
|
||||
var killIDMin uint64
|
||||
var killIDMax uint64
|
||||
var killAddr string
|
||||
var skipMe = true
|
||||
var maxAgeMs int64
|
||||
|
||||
for i := 0; i < len(args); i++ {
|
||||
arg := strings.ToUpper(string(args[i]))
|
||||
@@ -361,6 +382,12 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
|
||||
case "ID":
|
||||
if i+2 < len(args) {
|
||||
// ID min max format - kill all IDs in range
|
||||
if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killIDMin); err != nil {
|
||||
return 0, fmt.Errorf("invalid client ID: %s", args[i+1])
|
||||
}
|
||||
if _, err := fmt.Sscanf(string(args[i+2]), "%d", &killIDMax); err != nil {
|
||||
return 0, fmt.Errorf("invalid client ID: %s", args[i+2])
|
||||
}
|
||||
i += 2
|
||||
} else if i+1 < len(args) {
|
||||
// Single ID
|
||||
@@ -374,14 +401,35 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
|
||||
killAddr = string(args[i+1])
|
||||
i++
|
||||
}
|
||||
case "LADDR":
|
||||
// LADDR not fully supported, just skip
|
||||
if i+1 < len(args) {
|
||||
i++
|
||||
}
|
||||
case "SKIPME":
|
||||
if i+1 < len(args) {
|
||||
skipMe = strings.ToUpper(string(args[i+1])) == "YES"
|
||||
i++
|
||||
}
|
||||
case "MAXAGE":
|
||||
if i+1 < len(args) {
|
||||
if _, err := fmt.Sscanf(string(args[i+1]), "%d", &maxAgeMs); err != nil {
|
||||
return 0, fmt.Errorf("invalid maxage: %s", args[i+1])
|
||||
}
|
||||
if maxAgeMs <= 0 {
|
||||
return 0, fmt.Errorf("maxage must be greater than 0")
|
||||
}
|
||||
i++
|
||||
}
|
||||
case "TYPE":
|
||||
// TYPE filter handled separately in handleClient
|
||||
if i+1 < len(args) {
|
||||
i++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
for _, c := range clients {
|
||||
// Skip current client if requested
|
||||
if skipMe && c.ID() == currentClient.ID() {
|
||||
@@ -390,11 +438,29 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
|
||||
|
||||
shouldKill := false
|
||||
|
||||
if killID > 0 && c.ID() == killID {
|
||||
shouldKill = true
|
||||
// Check ID (single or range)
|
||||
if killID > 0 {
|
||||
shouldKill = c.ID() == killID
|
||||
} else if killIDMin > 0 || killIDMax > 0 {
|
||||
id := c.ID()
|
||||
if killIDMin > 0 && id < killIDMin {
|
||||
shouldKill = false
|
||||
} else if killIDMax > 0 && id > killIDMax {
|
||||
shouldKill = false
|
||||
} else {
|
||||
shouldKill = true
|
||||
}
|
||||
}
|
||||
if killAddr != "" && c.RemoteAddr() == killAddr {
|
||||
shouldKill = true
|
||||
|
||||
// Check ADDR
|
||||
if !shouldKill && killAddr != "" {
|
||||
shouldKill = c.RemoteAddr() == killAddr
|
||||
}
|
||||
|
||||
// Check MAXAGE - kill clients that have been connected longer than maxAge
|
||||
if shouldKill && maxAgeMs > 0 {
|
||||
age := now.Sub(c.CreatedAt()).Milliseconds()
|
||||
shouldKill = age <= maxAgeMs
|
||||
}
|
||||
|
||||
if shouldKill {
|
||||
@@ -406,9 +472,137 @@ func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (
|
||||
}
|
||||
}
|
||||
|
||||
// For single ID/ADDR that wasn't found, Redis returns error
|
||||
if killed == 0 && (killID > 0 || killAddr != "") {
|
||||
return 0, fmt.Errorf("No such client")
|
||||
}
|
||||
|
||||
return killed, nil
|
||||
}
|
||||
|
||||
// handleSelect processes SELECT command.
|
||||
// SELECT db - Switch to a different database (0-15).
|
||||
// For pekoe (pubsub-only), this only tracks the selected DB for client compatibility.
|
||||
// Redis default max DB is 16 (0-15), configurable via max-databases in config.
|
||||
func (s *Server) handleSelect(client *pubsub.Client, args [][]byte, writer *ConnWriter) error {
|
||||
if len(args) != 1 {
|
||||
writer.WriteError("ERR wrong number of arguments for 'select' command")
|
||||
return nil
|
||||
}
|
||||
|
||||
dbIndexStr := strings.TrimSpace(string(args[0]))
|
||||
dbIndex, err := strconv.Atoi(dbIndexStr)
|
||||
if err != nil {
|
||||
writer.WriteError("ERR invalid DB index")
|
||||
return nil
|
||||
}
|
||||
|
||||
if dbIndex < 0 || dbIndex >= MaxDatabases {
|
||||
writer.WriteError("ERR DB index is out of range")
|
||||
return nil
|
||||
}
|
||||
|
||||
client.SetSelectedDB(dbIndex)
|
||||
writer.WriteStatus("OK")
|
||||
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 {
|
||||
@@ -421,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 {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"math/rand"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
@@ -21,21 +22,37 @@ import (
|
||||
"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
|
||||
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.
|
||||
@@ -63,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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -103,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)
|
||||
}
|
||||
@@ -276,6 +301,10 @@ func (s *Server) processCommand(client *pubsub.Client, cmd *Command, writer *Con
|
||||
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) {
|
||||
@@ -311,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()
|
||||
|
||||
Reference in New Issue
Block a user