Initial commit

This commit is contained in:
eleksir
2026-04-02 23:56:28 +03:00
commit 9e30243533
34 changed files with 8413 additions and 0 deletions

View File

@@ -0,0 +1,959 @@
package compatibility
import (
"context"
"fmt"
"sync"
"testing"
"time"
"github.com/redis/go-redis/v9"
)
// TestGoRedisBasicSubscribe tests basic SUBSCRIBE functionality
func TestGoRedisBasicSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{
Addr: ts.Addr(),
})
defer func() { _ = rdb.Close() }()
pubsub := rdb.Subscribe(ctx, "test-channel")
defer func() { _ = pubsub.Close() }()
// Channel() returns a channel that receives messages (not subscription confirmations)
// The subscription confirmation is handled internally by go-redis
// We just need to verify we can get the channel without errors
ch := pubsub.Channel()
// Publish a message to verify subscription works
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "test-channel", "test-msg").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Now we should receive the message
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for message")
}
if msg.Channel != "test-channel" {
t.Errorf("Expected channel 'test-channel', got '%s'", msg.Channel)
}
if msg.Payload != "test-msg" {
t.Errorf("Expected 'test-msg', got '%s'", msg.Payload)
}
}
// TestGoRedisBasicPublish tests basic PUBLISH functionality
func TestGoRedisBasicPublish(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
pubsub := sub.Subscribe(ctx, "test-pub")
defer func() { _ = pubsub.Close() }()
ch := pubsub.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Publish message
err := pub.Publish(ctx, "test-pub", "hello world").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive message
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for message")
}
if msg.Payload != "hello world" {
t.Errorf("Expected 'hello world', got '%s'", msg.Payload)
}
if msg.Channel != "test-pub" {
t.Errorf("Expected channel 'test-pub', got '%s'", msg.Channel)
}
}
// TestGoRedisMultipleSubscribers tests multiple clients subscribing to same channel
func TestGoRedisMultipleSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
numSubscribers := 5
// Create multiple subscribers
var subscribers []*redis.Client
var pubsubs []*redis.PubSub
var channels []<-chan *redis.Message
for i := 0; i < numSubscribers; i++ {
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
subscribers = append(subscribers, rdb)
ps := rdb.Subscribe(ctx, "multi-sub")
pubsubs = append(pubsubs, ps)
channels = append(channels, ps.Channel())
}
defer func() {
for _, ps := range pubsubs {
_ = ps.Close()
}
for _, rdb := range subscribers {
_ = rdb.Close()
}
}()
// Give subscriptions time to register
time.Sleep(100 * time.Millisecond)
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Publish message
err := pub.Publish(ctx, "multi-sub", "broadcast").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// All subscribers should receive
for i, ch := range channels {
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Errorf("Subscriber %d timeout", i)
continue
}
if msg.Payload != "broadcast" {
t.Errorf("Subscriber %d got wrong payload: %s", i, msg.Payload)
}
}
}
// TestGoRedisUnsubscribe tests UNSUBSCRIBE functionality
func TestGoRedisUnsubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
pubsub := rdb.Subscribe(ctx, "unsub-test")
defer func() { _ = pubsub.Close() }()
// Verify subscription by publishing and receiving
ch := pubsub.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "unsub-test", "msg").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for message")
}
// Unsubscribe
err = pubsub.Unsubscribe(ctx, "unsub-test")
if err != nil {
t.Fatalf("Unsubscribe failed: %v", err)
}
// After unsubscribe, should not receive messages
// Give time for unsubscribe to process
time.Sleep(100 * time.Millisecond)
}
// TestGoRedisPatternSubscribe tests PSUBSCRIBE functionality
func TestGoRedisPatternSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
pubsub := rdb.PSubscribe(ctx, "news.*")
defer func() { _ = pubsub.Close() }()
// Get the channel for receiving messages
ch := pubsub.Channel()
// Publish to a matching channel to verify subscription
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "news.sports", "game results").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Should receive the message
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for pattern message")
}
if msg.Payload != "game results" {
t.Errorf("Expected 'game results', got '%s'", msg.Payload)
}
}
// TestGoRedisPatternMatch tests pattern matching with PUBLISH
func TestGoRedisPatternMatch(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber with pattern
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
pubsub := sub.PSubscribe(ctx, "news.*")
defer func() { _ = pubsub.Close() }()
ch := pubsub.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Publish to matching channel
err := pub.Publish(ctx, "news.sports", "game results").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive pattern message
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for pattern message")
}
if msg.Payload != "game results" {
t.Errorf("Expected 'game results', got '%s'", msg.Payload)
}
if msg.Channel != "news.sports" {
t.Errorf("Expected channel 'news.sports', got '%s'", msg.Channel)
}
}
// TestGoRedisMultiplePatterns tests subscribing to multiple patterns
func TestGoRedisMultiplePatterns(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
pubsub := rdb.PSubscribe(ctx, "news.*", "sports.*", "weather.*")
defer func() { _ = pubsub.Close() }()
// Get channel and verify subscription by publishing
ch := pubsub.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Publish to each pattern
_ = pub.Publish(ctx, "news.test", "news-msg").Err()
_ = pub.Publish(ctx, "sports.test", "sports-msg").Err()
_ = pub.Publish(ctx, "weather.test", "weather-msg").Err()
// Should receive 3 messages
for i := 0; i < 3; i++ {
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatalf("Timeout waiting for message %d", i)
}
}
}
// TestGoRedisMultipleChannels tests subscribing to multiple channels
func TestGoRedisMultipleChannels(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
pubsub := rdb.Subscribe(ctx, "ch1", "ch2", "ch3")
defer func() { _ = pubsub.Close() }()
// Verify by publishing to each channel
ch := pubsub.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
_ = pub.Publish(ctx, "ch1", "msg1").Err()
_ = pub.Publish(ctx, "ch2", "msg2").Err()
_ = pub.Publish(ctx, "ch3", "msg3").Err()
// Should receive 3 messages
for i := 0; i < 3; i++ {
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatalf("Timeout waiting for channel %d message", i)
}
}
}
// TestGoRedisPing tests PING command in non-pubsub mode
func TestGoRedisPing(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
result, err := rdb.Ping(ctx).Result()
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "PONG" {
t.Errorf("Expected 'PONG', got '%s'", result)
}
}
// TestGoRedisPingWithMessage tests PING with message in non-pubsub mode
func TestGoRedisPingWithMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
result, err := rdb.Do(ctx, "PING", "hello").Text()
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "hello" {
t.Errorf("Expected 'hello', got '%s'", result)
}
}
// TestGoRedisPubSubNumsub tests PUBSUB NUMSUB command
func TestGoRedisPubSubNumsub(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Create subscribers
sub1 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub1.Close() }()
ps1 := sub1.Subscribe(ctx, "channel1")
defer func() { _ = ps1.Close() }()
sub2 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub2.Close() }()
ps2 := sub2.Subscribe(ctx, "channel1", "channel2")
defer func() { _ = ps2.Close() }()
// Wait for subscriptions
time.Sleep(100 * time.Millisecond)
// Use regular client for PUBSUB command
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
result, err := pub.Do(ctx, "PUBSUB", "NUMSUB", "channel1", "channel2").Result()
if err != nil {
t.Fatalf("PUBSUB NUMSUB failed: %v", err)
}
// Result should be an array with channel,count pairs
t.Logf("PUBSUB NUMSUB result: %v", result)
}
// TestGoRedisPubSubChannels tests PUBSUB CHANNELS command
func TestGoRedisPubSubChannels(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscribe to channels
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
ps := sub.Subscribe(ctx, "test-cha-1", "test-cha-2")
defer func() { _ = ps.Close() }()
time.Sleep(100 * time.Millisecond)
// Query channels
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
result, err := pub.Do(ctx, "PUBSUB", "CHANNELS").Result()
if err != nil {
t.Fatalf("PUBSUB CHANNELS failed: %v", err)
}
t.Logf("PUBSUB CHANNELS result: %v", result)
}
// TestGoRedisPubSubNumpat tests PUBSUB NUMPAT command
func TestGoRedisPubSubNumpat(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Create pattern subscribers
sub1 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub1.Close() }()
ps1 := sub1.PSubscribe(ctx, "foo.*")
defer func() { _ = ps1.Close() }()
sub2 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub2.Close() }()
ps2 := sub2.PSubscribe(ctx, "bar.*")
defer func() { _ = ps2.Close() }()
time.Sleep(100 * time.Millisecond)
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
result, err := pub.Do(ctx, "PUBSUB", "NUMPAT").Result()
if err != nil {
t.Fatalf("PUBSUB NUMPAT failed: %v", err)
}
t.Logf("PUBSUB NUMPAT result: %v", result)
}
// TestGoRedisBinarySafe tests binary-safe message passing
func TestGoRedisBinarySafe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
pubsub := sub.Subscribe(ctx, "binary-test")
defer func() { _ = pubsub.Close() }()
ch := pubsub.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Binary message with null bytes and special chars
binaryMsg := []byte{0x00, 0x01, 0x02, 0xFF, 0xFE, '\r', '\n', 0x00}
err := pub.Publish(ctx, "binary-test", string(binaryMsg)).Err()
if err != nil {
t.Fatalf("Publish binary failed: %v", err)
}
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for binary message")
}
if msg.Payload != string(binaryMsg) {
t.Errorf("Binary message mismatch: got %v, want %v", []byte(msg.Payload), binaryMsg)
}
}
// TestGoRedisEmptyMessage tests publishing empty message
func TestGoRedisEmptyMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
pubsub := sub.Subscribe(ctx, "empty-test")
defer func() { _ = pubsub.Close() }()
ch := pubsub.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "empty-test", "").Err()
if err != nil {
t.Fatalf("Publish empty failed: %v", err)
}
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for empty message")
}
if msg.Payload != "" {
t.Errorf("Expected empty payload, got '%s'", msg.Payload)
}
}
// TestGoRedisLongMessage tests publishing long message
func TestGoRedisLongMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
pubsub := sub.Subscribe(ctx, "long-test")
defer func() { _ = pubsub.Close() }()
ch := pubsub.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Create 1MB message
longMsg := make([]byte, 1024*1024)
for i := range longMsg {
longMsg[i] = byte(i % 256)
}
err := pub.Publish(ctx, "long-test", string(longMsg)).Err()
if err != nil {
t.Fatalf("Publish long failed: %v", err)
}
msg, ok := waitForMessage(ch, 5*time.Second)
if !ok {
t.Fatal("Timeout waiting for long message")
}
if len(msg.Payload) != len(longMsg) {
t.Errorf("Message length mismatch: got %d, want %d", len(msg.Payload), len(longMsg))
}
}
// TestGoRedisConcurrentSubscribers tests concurrent subscription/unsubscription
func TestGoRedisConcurrentSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
var wg sync.WaitGroup
// Concurrent subscribe/unsubscribe
for i := 0; i < 20; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
ps := rdb.Subscribe(ctx, fmt.Sprintf("concurrent-%d", id%5))
defer func() { _ = ps.Close() }()
time.Sleep(50 * time.Millisecond)
}(i)
}
wg.Wait()
}
// TestGoRedisChannelWithSpecialChars tests channel names with special characters
func TestGoRedisChannelWithSpecialChars(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
channels := []string{
"channel-with-dash",
"channel_with_underscore",
"channel.with.dots",
"channel:with:colons",
"channel/with/slash",
}
for _, channelName := range channels {
t.Run(channelName, func(t *testing.T) {
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
ps := sub.Subscribe(ctx, channelName)
defer func() { _ = ps.Close() }()
ch := ps.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, channelName, "test").Err()
if err != nil {
t.Fatalf("Publish to %s failed: %v", channelName, err)
}
msg, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatalf("Timeout waiting for message on %s", channelName)
}
if msg.Channel != channelName {
t.Errorf("Channel mismatch: got %s, want %s", msg.Channel, channelName)
}
})
}
}
// TestGoRedisUnicodeMessages tests Unicode message content
func TestGoRedisUnicodeMessages(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
messages := []string{
"Hello, 世界!",
"Emoji test: 🎉🚀✨",
"Cyrillic: Привет мир",
"Arabic: مرحبا بالعالم",
"Japanese: こんにちは世界",
}
for _, msg := range messages {
t.Run(msg[:min(20, len(msg))], func(t *testing.T) {
// Subscriber
sub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub.Close() }()
ps := sub.Subscribe(ctx, "unicode-test")
defer func() { _ = ps.Close() }()
ch := ps.Channel()
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "unicode-test", msg).Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
received, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout waiting for Unicode message")
}
if received.Payload != msg {
t.Errorf("Unicode mismatch: got %q, want %q", received.Payload, msg)
}
})
}
}
// TestGoRedisPublishNoSubscribers tests PUBLISH with no subscribers
func TestGoRedisPublishNoSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
count, err := rdb.Publish(ctx, "no-subs-channel", "message").Result()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
if count != 0 {
t.Errorf("Expected 0 receivers, got %d", count)
}
}
// TestGoRedisSubscribeReturnCount tests that SUBSCRIBE returns correct count
func TestGoRedisSubscribeReturnCount(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
ps := rdb.Subscribe(ctx, "ch1")
defer func() { _ = ps.Close() }()
// The subscription is confirmed internally by go-redis
// Verify by publishing a message
ch := ps.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
err := pub.Publish(ctx, "ch1", "test").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout")
}
}
// TestGoRedisConcurrentPublishSubscribe tests concurrent pub/sub operations
func TestGoRedisConcurrentPublishSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
numMessages := 100
numSubscribers := 5
var mu sync.Mutex
receivedCounts := make([]int, numSubscribers)
var wg sync.WaitGroup
// Start subscribers
for i := 0; i < numSubscribers; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
ps := rdb.Subscribe(ctx, "concurrent-pub-test")
defer func() { _ = ps.Close() }()
ch := ps.Channel()
for j := 0; j < numMessages; j++ {
_, ok := waitForMessage(ch, 5*time.Second)
if ok {
mu.Lock()
receivedCounts[id]++
mu.Unlock()
}
}
}(i)
}
// Give subscribers time to connect
time.Sleep(200 * time.Millisecond)
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
for i := 0; i < numMessages; i++ {
err := pub.Publish(ctx, "concurrent-pub-test", fmt.Sprintf("msg-%d", i)).Err()
if err != nil {
t.Errorf("Publish %d failed: %v", i, err)
}
}
wg.Wait()
// Verify all subscribers received all messages
totalReceived := 0
for i, count := range receivedCounts {
totalReceived += count
if count < numMessages-1 { // Allow 1 message loss due to timing
t.Logf("Subscriber %d received %d/%d messages", i, count, numMessages)
}
}
t.Logf("Total messages received: %d/%d", totalReceived, numMessages*numSubscribers)
}
// TestGoRedisSubscribePatternThenChannel tests mixing pattern and channel subscriptions
func TestGoRedisSubscribePatternThenChannel(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
// Subscribe to pattern first
ps := rdb.PSubscribe(ctx, "test.*")
defer func() { _ = ps.Close() }()
ch := ps.Channel()
// Subscribe to a regular channel
err := ps.Subscribe(ctx, "direct-channel")
if err != nil {
t.Fatalf("Subscribe to channel failed: %v", err)
}
// Verify both subscriptions work
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
_ = pub.Publish(ctx, "test.foo", "pattern-msg").Err()
_ = pub.Publish(ctx, "direct-channel", "channel-msg").Err()
// Should receive both messages
for i := 0; i < 2; i++ {
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatalf("Timeout waiting for message %d", i)
}
}
}
// TestGoRedisUnsubscribeNoChannel tests UNSUBSCRIBE when not subscribed
func TestGoRedisUnsubscribeNoChannel(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
ps := rdb.Subscribe(ctx, "temp-channel")
defer func() { _ = ps.Close() }()
// Verify subscription works
ch := ps.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
_ = pub.Publish(ctx, "temp-channel", "test").Err()
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout")
}
// Unsubscribe
err := ps.Unsubscribe(ctx, "temp-channel")
if err != nil {
t.Fatalf("Unsubscribe failed: %v", err)
}
}
// TestGoRedisPunsubscribe tests PUNSUBSCRIBE functionality
func TestGoRedisPunsubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
ps := rdb.PSubscribe(ctx, "pattern.*")
defer func() { _ = ps.Close() }()
// Verify subscription
ch := ps.Channel()
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
_ = pub.Publish(ctx, "pattern.test", "msg").Err()
_, ok := waitForMessage(ch, 2*time.Second)
if !ok {
t.Fatal("Timeout")
}
// Punsubscribe
err := ps.PUnsubscribe(ctx, "pattern.*")
if err != nil {
t.Fatalf("PUnsubscribe failed: %v", err)
}
}
// TestGoRedisMixedPubSub tests a mix of channel and pattern subscriptions
func TestGoRedisMixedPubSub(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
// Subscriber 1: regular channel
sub1 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub1.Close() }()
ps1 := sub1.Subscribe(ctx, "mixed-channel")
defer func() { _ = ps1.Close() }()
// Subscriber 2: pattern
sub2 := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = sub2.Close() }()
ps2 := sub2.PSubscribe(ctx, "mixed*")
defer func() { _ = ps2.Close() }()
// Wait for subscriptions
time.Sleep(100 * time.Millisecond)
// Publisher
pub := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = pub.Close() }()
// Publish to "mixed-channel" - both should receive
err := pub.Publish(ctx, "mixed-channel", "hello").Err()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
ch1 := ps1.Channel()
ch2 := ps2.Channel()
msg1, ok := waitForMessage(ch1, 2*time.Second)
if !ok {
t.Error("Channel subscriber timeout")
} else if msg1.Payload != "hello" {
t.Errorf("Wrong payload: %s", msg1.Payload)
}
msg2, ok := waitForMessage(ch2, 2*time.Second)
if !ok {
t.Error("Pattern subscriber timeout")
} else if msg2.Payload != "hello" {
t.Errorf("Wrong payload: %s", msg2.Payload)
}
}
// TestGoRedisQuit tests QUIT command properly closes connection
func TestGoRedisQuit(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
rdb := redis.NewClient(&redis.Options{Addr: ts.Addr()})
defer func() { _ = rdb.Close() }()
// Use Do to send QUIT command directly (Quit method may not be implemented in all versions)
result, err := rdb.Do(ctx, "QUIT").Text()
if err != nil {
t.Fatalf("Quit failed: %v", err)
}
if result != "OK" {
t.Errorf("Expected 'OK', got '%s'", result)
}
}
func min(a, b int) int {
if a < b {
return a
}
return b
}

View File

@@ -0,0 +1,118 @@
package compatibility
import (
"context"
"net"
"sync"
"testing"
"time"
"github.com/eleksir/pekoe/internal/config"
"github.com/eleksir/pekoe/internal/pubsub"
"github.com/eleksir/pekoe/internal/server"
)
// testServer manages a test server instance for compatibility tests
type testServer struct {
cfg *config.Config
server *server.Server
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
addr string
}
// startTestServer starts a pekoe server on a random available port for testing
func startTestServer(t *testing.T) *testServer {
t.Helper()
// Find an available port
port := findAvailablePort(t)
cfg := &config.Config{
Bind: "127.0.0.1",
Port: port,
MaxClients: 100,
ReadTimeout: 0,
WriteTimeout: 0,
MaxMsgBytes: 10 * 1024 * 1024, // 10MB
LogLevel: "error", // Suppress logs in tests
}
cfg.SetDefaults()
hub := pubsub.NewHub()
srv := server.New(cfg, hub)
ctx, cancel := context.WithCancel(context.Background())
ts := &testServer{
cfg: cfg,
server: srv,
ctx: ctx,
cancel: cancel,
addr: cfg.Addr(),
}
// Start server in background
ts.wg.Add(1)
go func() {
defer ts.wg.Done()
_ = srv.Run()
}()
// Wait for server to be ready
ts.waitForReady(t)
return ts
}
// waitForReady blocks until the server accepts connections
func (ts *testServer) waitForReady(t *testing.T) {
t.Helper()
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
conn, err := net.DialTimeout("tcp", ts.addr, 100*time.Millisecond)
if err == nil {
_ = conn.Close()
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("Server at %s not ready after 5 seconds", ts.addr)
}
// Addr returns the server address for client connections
func (ts *testServer) Addr() string {
return ts.addr
}
// Close shuts down the test server
func (ts *testServer) Close() {
ts.cancel()
// Give time for graceful shutdown
time.Sleep(100 * time.Millisecond)
}
// findAvailablePort finds an available TCP port for testing
func findAvailablePort(t *testing.T) int {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("Failed to find available port: %v", err)
}
port := listener.Addr().(*net.TCPAddr).Port
_ = listener.Close()
return port
}
// waitForMessage waits for a message on a channel with timeout
func waitForMessage[T any](ch <-chan T, timeout time.Duration) (T, bool) {
select {
case msg := <-ch:
return msg, true
case <-time.After(timeout):
var zero T
return zero, false
}
}

View File

@@ -0,0 +1,833 @@
package compatibility
import (
"fmt"
"sync"
"testing"
"time"
"github.com/gomodule/redigo/redis"
)
// dialRedis creates a new redigo connection for testing
func dialRedis(t *testing.T, addr string) redis.Conn {
t.Helper()
conn, err := redis.Dial("tcp", addr,
redis.DialConnectTimeout(5*time.Second),
redis.DialReadTimeout(5*time.Second),
redis.DialWriteTimeout(5*time.Second),
)
if err != nil {
t.Fatalf("Failed to dial: %v", err)
}
return conn
}
// TestRedigoBasicSubscribe tests basic SUBSCRIBE functionality
func TestRedigoBasicSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
// Send SUBSCRIBE
err := conn.Send("SUBSCRIBE", "rigo-test-channel")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = conn.Flush()
// Read subscription confirmation
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
// Should be ["subscribe", "channel", count]
if len(msg) != 3 {
t.Fatalf("Expected 3 elements, got %d", len(msg))
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "subscribe" {
t.Errorf("Expected 'subscribe', got '%s'", msgType)
}
channel, _ := redis.String(msg[1], nil)
if channel != "rigo-test-channel" {
t.Errorf("Expected 'rigo-test-channel', got '%s'", channel)
}
}
// TestRedigoBasicPublish tests basic PUBLISH functionality
func TestRedigoBasicPublish(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-pub-test")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
// Consume subscription confirmation
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive subscription failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
reply, err := redis.Int(pubConn.Do("PUBLISH", "rigo-pub-test", "hello redigo"))
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
t.Logf("Published to %d subscribers", reply)
// Receive message
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive message failed: %v", err)
}
// Should be ["message", "channel", payload]
if len(msg) != 3 {
t.Fatalf("Expected 3 elements, got %d", len(msg))
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "message" {
t.Errorf("Expected 'message', got '%s'", msgType)
}
payload, _ := redis.String(msg[2], nil)
if payload != "hello redigo" {
t.Errorf("Expected 'hello redigo', got '%s'", payload)
}
}
// TestRedigoMultipleSubscribers tests multiple subscribers
func TestRedigoMultipleSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
numSubs := 3
conns := make([]redis.Conn, numSubs)
// Create subscribers
for i := 0; i < numSubs; i++ {
conn := dialRedis(t, ts.Addr())
conns[i] = conn
err := conn.Send("SUBSCRIBE", "rigo-multi")
if err != nil {
t.Fatalf("Send failed for %d: %v", i, err)
}
_ = conn.Flush()
// Consume confirmation
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed for %d: %v", i, err)
}
}
defer func() {
for _, c := range conns {
_ = c.Close()
}
}()
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
_, err := pubConn.Do("PUBLISH", "rigo-multi", "broadcast")
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// All should receive
for i, conn := range conns {
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Errorf("Subscriber %d failed to receive: %v", i, err)
continue
}
payload, _ := redis.String(msg[2], nil)
if payload != "broadcast" {
t.Errorf("Subscriber %d got wrong payload: %s", i, payload)
}
}
}
// TestRedigoUnsubscribe tests UNSUBSCRIBE functionality
func TestRedigoUnsubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
// Subscribe
err := conn.Send("SUBSCRIBE", "rigo-unsub-test")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = conn.Flush()
// Consume confirmation
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
// Unsubscribe
err = conn.Send("UNSUBSCRIBE", "rigo-unsub-test")
if err != nil {
t.Fatalf("Send unsubscribe failed: %v", err)
}
_ = conn.Flush()
// Receive unsubscribe confirmation
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive unsubscribe confirmation failed: %v", err)
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "unsubscribe" {
t.Errorf("Expected 'unsubscribe', got '%s'", msgType)
}
}
// TestRedigoPatternSubscribe tests PSUBSCRIBE functionality
func TestRedigoPatternSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
// PSubscribe
err := conn.Send("PSUBSCRIBE", "rigo-news.*")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = conn.Flush()
// Receive confirmation
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "psubscribe" {
t.Errorf("Expected 'psubscribe', got '%s'", msgType)
}
pattern, _ := redis.String(msg[1], nil)
if pattern != "rigo-news.*" {
t.Errorf("Expected 'rigo-news.*', got '%s'", pattern)
}
}
// TestRedigoPatternMatch tests pattern matching with PUBLISH
func TestRedigoPatternMatch(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
// Subscriber with pattern
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("PSUBSCRIBE", "rigo-pattern.*")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
// Consume confirmation
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
t.Log("Pattern subscription confirmed, now publishing...")
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
reply, err := redis.Int(pubConn.Do("PUBLISH", "rigo-pattern.match", "pattern message"))
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
t.Logf("Publish returned: %d", reply)
// Receive pmessage
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
// Should be ["pmessage", "pattern", "channel", "payload"]
if len(msg) != 4 {
t.Fatalf("Expected 4 elements, got %d", len(msg))
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "pmessage" {
t.Errorf("Expected 'pmessage', got '%s'", msgType)
}
pattern, _ := redis.String(msg[1], nil)
if pattern != "rigo-pattern.*" {
t.Errorf("Expected pattern 'rigo-pattern.*', got '%s'", pattern)
}
channel, _ := redis.String(msg[2], nil)
if channel != "rigo-pattern.match" {
t.Errorf("Expected channel 'rigo-pattern.match', got '%s'", channel)
}
payload, _ := redis.String(msg[3], nil)
if payload != "pattern message" {
t.Errorf("Expected 'pattern message', got '%s'", payload)
}
}
// TestRedigoMultiplePatterns tests subscribing to multiple patterns
func TestRedigoMultiplePatterns(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
patterns := []string{"rigo-pat1.*", "rigo-pat2.*", "rigo-pat3.*"}
for _, pattern := range patterns {
err := conn.Send("PSUBSCRIBE", pattern)
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = conn.Flush()
// Consume confirmation
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
}
}
// TestRedigoPing tests PING command
func TestRedigoPing(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
result, err := redis.String(conn.Do("PING"))
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "PONG" {
t.Errorf("Expected 'PONG', got '%s'", result)
}
}
// TestRedigoPingWithMessage tests PING with message
func TestRedigoPingWithMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
result, err := redis.String(conn.Do("PING", "hello"))
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "hello" {
t.Errorf("Expected 'hello', got '%s'", result)
}
}
// TestRedigoBinarySafe tests binary-safe messages
func TestRedigoBinarySafe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-binary")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
binaryMsg := []byte{0x00, 0x01, 0x02, 0xFF, 0xFE, '\r', '\n', 0x00}
_, err = pubConn.Do("PUBLISH", "rigo-binary", binaryMsg)
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive message
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
payload, _ := redis.String(msg[2], nil)
if []byte(payload)[0] != binaryMsg[0] {
t.Errorf("Binary message mismatch")
}
}
// TestRedigisEmptyMessage tests empty message
func TestRedigisEmptyMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-empty")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
_, err = pubConn.Do("PUBLISH", "rigo-empty", "")
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive message
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
payload, _ := redis.String(msg[2], nil)
if payload != "" {
t.Errorf("Expected empty payload, got '%s'", payload)
}
}
// TestRedigoLongMessage tests long message
func TestRedigoLongMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-long")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
// 100KB message
longMsg := make([]byte, 100*1024)
for i := range longMsg {
longMsg[i] = byte(i % 256)
}
_, err = pubConn.Do("PUBLISH", "rigo-long", longMsg)
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive message
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
payload, _ := redis.String(msg[2], nil)
if len(payload) != len(longMsg) {
t.Errorf("Message length mismatch: got %d, want %d", len(payload), len(longMsg))
}
}
// TestRedigoConcurrentSubscribers tests concurrent operations
func TestRedigoConcurrentSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
var wg sync.WaitGroup
for i := 0; i < 10; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
channel := fmt.Sprintf("rigo-concurrent-%d", id%3)
err := conn.Send("SUBSCRIBE", channel)
if err != nil {
t.Errorf("Send failed: %v", err)
return
}
_ = conn.Flush()
_, err = redis.Values(conn.Receive())
if err != nil {
t.Errorf("Receive failed: %v", err)
return
}
time.Sleep(50 * time.Millisecond)
}(i)
}
wg.Wait()
}
// TestRedigoPublishNoSubscribers tests PUBLISH with no subscribers
func TestRedigoPublishNoSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
count, err := redis.Int(conn.Do("PUBLISH", "rigo-no-subs", "message"))
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
if count != 0 {
t.Errorf("Expected 0 receivers, got %d", count)
}
}
// TestRedigoUnicodeMessages tests Unicode messages
func TestRedigoUnicodeMessages(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
messages := []string{
"Hello, 世界!",
"Emoji: 🎉",
"Cyrillic: Привет",
}
for _, testMsg := range messages {
t.Run(testMsg[:min(20, len(testMsg))], func(t *testing.T) {
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-unicode")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
_, err = pubConn.Do("PUBLISH", "rigo-unicode", testMsg)
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
payload, _ := redis.String(msg[2], nil)
if payload != testMsg {
t.Errorf("Unicode mismatch: got %q, want %q", payload, testMsg)
}
})
}
}
// TestRedigoChannelWithSpecialChars tests channels with special characters
func TestRedigoChannelWithSpecialChars(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
channels := []string{
"rigo-with-dash",
"rigo_with_underscore",
"rigo.with.dots",
"rigo:with:colons",
}
for _, channelName := range channels {
t.Run(channelName, func(t *testing.T) {
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", channelName)
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
_, err = pubConn.Do("PUBLISH", channelName, "test")
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
// Receive
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
channel, _ := redis.String(msg[1], nil)
if channel != channelName {
t.Errorf("Channel mismatch: got %s, want %s", channel, channelName)
}
})
}
}
// TestRedigoMixedSubscribePSubscribe tests mixing SUBSCRIBE and PSUBSCRIBE
// nolint: cyclop
func TestRedigoMixedSubscribePSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
// Subscribe to channel
err := conn.Send("SUBSCRIBE", "rigo-mixed-ch")
if err != nil {
t.Fatalf("Send subscribe failed: %v", err)
}
_ = conn.Flush()
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive subscribe confirmation failed: %v", err)
}
// Also subscribe to pattern
err = conn.Send("PSUBSCRIBE", "rigo-mixed-*")
if err != nil {
t.Fatalf("Send psubscribe failed: %v", err)
}
_ = conn.Flush()
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive psubscribe confirmation failed: %v", err)
}
// Publish to channel - should receive via channel subscription
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
_, err = pubConn.Do("PUBLISH", "rigo-mixed-ch", "channel msg")
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "message" {
t.Errorf("Expected 'message', got '%s'", msgType)
}
// Publish to pattern-matching channel - should receive via pattern subscription
_, err = pubConn.Do("PUBLISH", "rigo-mixed-pattern", "pattern msg")
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
msg, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive failed: %v", err)
}
msgType, _ = redis.String(msg[0], nil)
if msgType != "pmessage" {
t.Errorf("Expected 'pmessage', got '%s'", msgType)
}
}
// TestRedigoPunsubscribe tests PUNSUBSCRIBE functionality
func TestRedigoPunsubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
conn := dialRedis(t, ts.Addr())
defer func() { _ = conn.Close() }()
// PSubscribe
err := conn.Send("PSUBSCRIBE", "rigo-punsub.*")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = conn.Flush()
_, err = redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Punsubscribe
err = conn.Send("PUNSUBSCRIBE", "rigo-punsub.*")
if err != nil {
t.Fatalf("Send punsubscribe failed: %v", err)
}
_ = conn.Flush()
msg, err := redis.Values(conn.Receive())
if err != nil {
t.Fatalf("Receive punsubscribe confirmation failed: %v", err)
}
msgType, _ := redis.String(msg[0], nil)
if msgType != "punsubscribe" {
t.Errorf("Expected 'punsubscribe', got '%s'", msgType)
}
}
// TestRedigoConcurrentPublishSubscribe tests concurrent pub/sub
func TestRedigoConcurrentPublishSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
numMessages := 50
// Subscriber
subConn := dialRedis(t, ts.Addr())
defer func() { _ = subConn.Close() }()
err := subConn.Send("SUBSCRIBE", "rigo-concurrent-pub")
if err != nil {
t.Fatalf("Send failed: %v", err)
}
_ = subConn.Flush()
_, err = redis.Values(subConn.Receive())
if err != nil {
t.Fatalf("Receive confirmation failed: %v", err)
}
// Publisher in goroutine
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
pubConn := dialRedis(t, ts.Addr())
defer func() { _ = pubConn.Close() }()
for i := 0; i < numMessages; i++ {
_, err := pubConn.Do("PUBLISH", "rigo-concurrent-pub", fmt.Sprintf("msg-%d", i))
if err != nil {
t.Errorf("Publish %d failed: %v", i, err)
return
}
time.Sleep(time.Millisecond)
}
}()
// Receive messages
received := 0
timeout := time.After(10 * time.Second)
for received < numMessages {
select {
case <-timeout:
t.Logf("Received %d/%d messages", received, numMessages)
wg.Wait()
return
default:
}
msg, err := redis.Values(subConn.Receive())
if err != nil {
t.Errorf("Receive failed: %v", err)
break
}
msgType, _ := redis.String(msg[0], nil)
if msgType == "message" {
received++
}
}
wg.Wait()
t.Logf("Received all %d messages", received)
}

View File

@@ -0,0 +1,217 @@
package compatibility
import (
"context"
"fmt"
"sync"
"testing"
"github.com/redis/rueidis"
)
// newRueidisClient creates a rueidis client with proper options for pekoe
func newRueidisClient(t *testing.T, addr string) rueidis.Client {
t.Helper()
client, err := rueidis.NewClient(rueidis.ClientOption{
InitAddress: []string{addr},
DisableCache: true,
})
if err != nil {
t.Fatalf("Failed to create rueidis client: %v", err)
}
return client
}
// TestRueidisPing tests PING command
func TestRueidisPing(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
result, err := client.Do(ctx, client.B().Ping().Build()).ToString()
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "PONG" {
t.Errorf("Expected 'PONG', got '%s'", result)
}
}
// TestRueidisPingWithMessage tests PING with message
func TestRueidisPingWithMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
result, err := client.Do(ctx, client.B().Ping().Message("hello").Build()).ToString()
if err != nil {
t.Fatalf("Ping failed: %v", err)
}
if result != "hello" {
t.Errorf("Expected 'hello', got '%s'", result)
}
}
// TestRueidisBasicPublish tests PUBLISH and verifies count
func TestRueidisBasicPublish(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
pubClient := newRueidisClient(t, ts.Addr())
defer pubClient.Close()
// Publish to channel with no subscribers
count, err := pubClient.Do(ctx, pubClient.B().Publish().Channel("r-test").Message("hello").Build()).AsInt64()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
if count != 0 {
t.Errorf("Expected 0 subscribers, got %d", count)
}
}
// TestRueidisPublishNoSubscribers tests PUBLISH with no subscribers
func TestRueidisPublishNoSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
n, err := client.Do(ctx, client.B().Publish().Channel("r-no-subs").Message("msg").Build()).AsInt64()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
if n != 0 {
t.Errorf("Expected 0 receivers, got %d", n)
}
}
// TestRueidisBinarySafe tests binary message handling
func TestRueidisBinarySafe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
binaryMsg := []byte{0x00, 0x01, 0x02, 0xFF, 0xFE}
// Just verify publish works with binary data
_, err := client.Do(ctx, client.B().Publish().Channel("r-binary").Message(string(binaryMsg)).Build()).AsInt64()
if err != nil {
t.Fatalf("Publish binary failed: %v", err)
}
}
// TestRueidisEmptyMessage tests empty message
func TestRueidisEmptyMessage(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
_, err := client.Do(ctx, client.B().Publish().Channel("r-empty").Message("").Build()).AsInt64()
if err != nil {
t.Fatalf("Publish empty failed: %v", err)
}
}
// TestRueidisConcurrentSubscribers tests concurrent client creation
func TestRueidisConcurrentSubscribers(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
var wg sync.WaitGroup
successCount := 0
var mu sync.Mutex
for i := 0; i < 10; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
// Just verify we can ping
ctx := context.Background()
_, err := client.Do(ctx, client.B().Ping().Build()).ToString()
if err != nil {
t.Errorf("Client %d ping failed: %v", id, err)
return
}
mu.Lock()
successCount++
mu.Unlock()
}(i)
}
wg.Wait()
if successCount != 10 {
t.Errorf("Only %d/10 clients succeeded", successCount)
}
}
// TestRueidisUnicodeMessages tests Unicode messages
func TestRueidisUnicodeMessages(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
client := newRueidisClient(t, ts.Addr())
defer client.Close()
messages := []string{
"Hello, 世界!",
"Emoji: 🎉",
"Cyrillic: Привет",
}
for _, testMsg := range messages {
t.Run(testMsg[:min(20, len(testMsg))], func(t *testing.T) {
_, err := client.Do(ctx, client.B().Publish().Channel("r-unicode").Message(testMsg).Build()).AsInt64()
if err != nil {
t.Fatalf("Publish failed: %v", err)
}
})
}
}
// TestRueidisConcurrentPublishSubscribe tests concurrent pub/sub operations
func TestRueidisConcurrentPublishSubscribe(t *testing.T) {
ts := startTestServer(t)
defer ts.Close()
ctx := context.Background()
numMessages := 50
pubClient := newRueidisClient(t, ts.Addr())
defer pubClient.Close()
for i := 0; i < numMessages; i++ {
_, err := pubClient.Do(ctx, pubClient.B().Publish().Channel("r-concurrent").Message(fmt.Sprintf("msg-%d", i)).Build()).AsInt64()
if err != nil {
t.Errorf("Publish %d failed: %v", i, err)
}
}
t.Logf("Published %d messages", numMessages)
}

169
internal/config/config.go Normal file
View File

@@ -0,0 +1,169 @@
// @tag config
// @tag core
package config
import (
"encoding/json"
"fmt"
"net"
"os"
"strconv"
"time"
"github.com/eleksir/pekoe/internal/logger"
)
// 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
)
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
}
// SetDefaults applies sensible defaults when config values are missing or invalid.
// This ensures the server can start even with incomplete configuration.
// nolint: cyclop
func (c *Config) SetDefaults() {
if c.Bind == "" {
c.Bind = DefaultBind
}
if c.Port <= 0 || c.Port > 65535 {
c.Port = DefaultPort
}
if c.MaxClients <= 0 {
c.MaxClients = DefaultMaxClients
}
if c.ReadTimeout < 0 {
c.ReadTimeout = 0
}
if c.WriteTimeout < 0 {
c.WriteTimeout = 0
}
if c.MaxMsgBytes <= 0 {
c.MaxMsgBytes = DefaultMaxMsgBytes
}
if c.LogLevel == "" {
c.LogLevel = DefaultLogLevel
}
if c.LogFile == "" {
c.LogFile = DefaultLogFile
}
if c.PidFile == "" {
c.PidFile = DefaultPidFile
}
}
// Validate checks that all configuration values are within acceptable bounds.
// Returns an error if validation fails, which should prevent server startup.
func (c *Config) Validate() error {
if c.Bind == "" {
return fmt.Errorf("bind address cannot be empty")
}
if net.ParseIP(c.Bind) == nil {
return fmt.Errorf("bind must be a valid IP address (e.g. 0.0.0.0, ::, ::1, 192.168.1.1)")
}
if c.Port <= 0 || c.Port > 65535 {
return fmt.Errorf("port must be between 1 and 65535")
}
if c.MaxClients <= 0 {
return fmt.Errorf("max_clients must be positive")
}
if c.MaxMsgBytes <= 0 {
return fmt.Errorf("max_message_bytes must be positive")
}
// Validate log level against allowed values
validLevels := map[string]bool{
"debug": true,
"info": true,
"warn": true,
"error": true,
}
if !validLevels[c.LogLevel] {
return fmt.Errorf("invalid log_level: %s (must be debug, info, warn, or error)", c.LogLevel)
}
return nil
}
// ReadTimeoutDuration converts the config's read timeout (in seconds) to a Duration.
// Zero means no timeout, allowing the server to block on reads indefinitely.
func (c *Config) ReadTimeoutDuration() time.Duration {
if c.ReadTimeout == 0 {
return 0
}
return time.Duration(c.ReadTimeout) * time.Second
}
// WriteTimeoutDuration converts the config's write timeout (in seconds) to a Duration.
func (c *Config) WriteTimeoutDuration() time.Duration {
if c.WriteTimeout == 0 {
return 0
}
return time.Duration(c.WriteTimeout) * time.Second
}
// Load reads and parses the configuration from a JSON file.
// It applies defaults for missing/invalid values and validates the final config.
func Load(path string) (*Config, error) {
//nolint: gosec
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read config file: %w", err)
}
var cfg Config
if err := json.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("failed to parse config file: %w", err)
}
// Apply defaults before validation to ensure minimum viable config
cfg.SetDefaults()
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("config validation failed: %w", err)
}
logger.SetLevel(cfg.LogLevel)
if cfg.LogFile != "" {
if err := logger.SetOutput(cfg.LogFile); err != nil {
return nil, fmt.Errorf("failed to set log file: %w", err)
}
}
return &cfg, nil
}
// Addr returns the bind address in host:port format for use in net.Listen.
// IPv6 addresses are wrapped in brackets: [::]:6379.
func (c *Config) Addr() string {
ip := net.ParseIP(c.Bind)
if ip != nil && ip.To4() == nil {
// IPv6 address — wrap in brackets
return "[" + c.Bind + "]:" + strconv.Itoa(c.Port)
}
return c.Bind + ":" + strconv.Itoa(c.Port)
}
// String returns a JSON representation of the config for debugging.
func (c *Config) String() string {
b, _ := json.MarshalIndent(c, "", " ")
return string(b)
}

View File

@@ -0,0 +1,435 @@
package config
import (
"os"
"testing"
)
// nolint: cyclop
func TestSetDefaults(t *testing.T) {
tests := []struct {
name string
cfg *Config
check func(*Config)
}{
{
name: "empty config gets all defaults",
cfg: &Config{},
check: func(c *Config) {
if c.Bind != DefaultBind {
t.Errorf("Bind = %q, want %q", c.Bind, DefaultBind)
}
if c.Port != DefaultPort {
t.Errorf("Port = %d, want %d", c.Port, DefaultPort)
}
if c.MaxClients != DefaultMaxClients {
t.Errorf("MaxClients = %d, want %d", c.MaxClients, DefaultMaxClients)
}
if c.MaxMsgBytes != DefaultMaxMsgBytes {
t.Errorf("MaxMsgBytes = %d, want %d", c.MaxMsgBytes, DefaultMaxMsgBytes)
}
if c.LogLevel != DefaultLogLevel {
t.Errorf("LogLevel = %q, want %q", c.LogLevel, DefaultLogLevel)
}
},
},
{
name: "port zero uses default",
cfg: &Config{Port: 0},
check: func(c *Config) {
if c.Port != DefaultPort {
t.Errorf("Port = %d, want %d", c.Port, DefaultPort)
}
},
},
{
name: "port out of range uses default",
cfg: &Config{Port: 70000},
check: func(c *Config) {
if c.Port != DefaultPort {
t.Errorf("Port = %d, want %d", c.Port, DefaultPort)
}
},
},
{
name: "negative port uses default",
cfg: &Config{Port: -1},
check: func(c *Config) {
if c.Port != DefaultPort {
t.Errorf("Port = %d, want %d", c.Port, DefaultPort)
}
},
},
{
name: "negative timeout becomes zero",
cfg: &Config{ReadTimeout: -5, WriteTimeout: -10},
check: func(c *Config) {
if c.ReadTimeout != 0 {
t.Errorf("ReadTimeout = %d, want 0", c.ReadTimeout)
}
if c.WriteTimeout != 0 {
t.Errorf("WriteTimeout = %d, want 0", c.WriteTimeout)
}
},
},
{
name: "zero max_clients uses default",
cfg: &Config{MaxClients: 0},
check: func(c *Config) {
if c.MaxClients != DefaultMaxClients {
t.Errorf("MaxClients = %d, want %d", c.MaxClients, DefaultMaxClients)
}
},
},
{
name: "zero max_message_bytes uses default",
cfg: &Config{MaxMsgBytes: 0},
check: func(c *Config) {
if c.MaxMsgBytes != DefaultMaxMsgBytes {
t.Errorf("MaxMsgBytes = %d, want %d", c.MaxMsgBytes, DefaultMaxMsgBytes)
}
},
},
{
name: "empty log_level uses default",
cfg: &Config{LogLevel: ""},
check: func(c *Config) {
if c.LogLevel != DefaultLogLevel {
t.Errorf("LogLevel = %q, want %q", c.LogLevel, DefaultLogLevel)
}
},
},
{
name: "valid values are preserved",
cfg: &Config{
Bind: "127.0.0.1",
Port: 6380,
MaxClients: 500,
ReadTimeout: 30,
WriteTimeout: 60,
MaxMsgBytes: 1024,
LogLevel: "debug",
},
check: func(c *Config) {
if c.Bind != "127.0.0.1" {
t.Errorf("Bind = %q, want 127.0.0.1", c.Bind)
}
if c.Port != 6380 {
t.Errorf("Port = %d, want 6380", c.Port)
}
if c.MaxClients != 500 {
t.Errorf("MaxClients = %d, want 500", c.MaxClients)
}
if c.ReadTimeout != 30 {
t.Errorf("ReadTimeout = %d, want 30", c.ReadTimeout)
}
if c.WriteTimeout != 60 {
t.Errorf("WriteTimeout = %d, want 60", c.WriteTimeout)
}
if c.MaxMsgBytes != 1024 {
t.Errorf("MaxMsgBytes = %d, want 1024", c.MaxMsgBytes)
}
if c.LogLevel != "debug" {
t.Errorf("LogLevel = %q, want debug", c.LogLevel)
}
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tt.cfg.SetDefaults()
tt.check(tt.cfg)
})
}
}
func TestValidate(t *testing.T) {
tests := []struct {
name string
cfg *Config
wantErr bool
errMsg string
}{
{
name: "valid config passes",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: false,
},
{
name: "valid IPv6 bind passes",
cfg: &Config{Bind: "::", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: false,
},
{
name: "valid IPv6 loopback passes",
cfg: &Config{Bind: "::1", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: false,
},
{
name: "empty bind fails",
cfg: &Config{Bind: "", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "bind address cannot be empty",
},
{
name: "invalid bind fails",
cfg: &Config{Bind: "not-an-ip", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "bind must be a valid IP address (e.g. 0.0.0.0, ::, ::1, 192.168.1.1)",
},
{
name: "port zero fails",
cfg: &Config{Bind: "0.0.0.0", Port: 0, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "port must be between 1 and 65535",
},
{
name: "port out of range fails",
cfg: &Config{Bind: "0.0.0.0", Port: 70000, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "port must be between 1 and 65535",
},
{
name: "negative port fails",
cfg: &Config{Bind: "0.0.0.0", Port: -1, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "port must be between 1 and 65535",
},
{
name: "zero max_clients fails",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: 0, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "max_clients must be positive",
},
{
name: "negative max_clients fails",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: -5, MaxMsgBytes: 1024, LogLevel: "info"},
wantErr: true,
errMsg: "max_clients must be positive",
},
{
name: "zero max_message_bytes fails",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: 100, MaxMsgBytes: 0, LogLevel: "info"},
wantErr: true,
errMsg: "max_message_bytes must be positive",
},
{
name: "invalid log_level fails",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "invalid"},
wantErr: true,
errMsg: "invalid log_level: invalid (must be debug, info, warn, or error)",
},
{
name: "valid log levels pass",
cfg: &Config{Bind: "0.0.0.0", Port: 6379, MaxClients: 100, MaxMsgBytes: 1024, LogLevel: "debug"},
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.cfg.Validate()
if tt.wantErr {
if err == nil {
t.Errorf("Validate() expected error, got nil")
return
}
if err.Error() != tt.errMsg {
t.Errorf("Validate() error = %q, want %q", err.Error(), tt.errMsg)
}
} else if err != nil {
t.Errorf("Validate() unexpected error: %v", err)
}
})
}
}
// nolint: cyclop
func TestLoad(t *testing.T) {
tmpfile, err := os.CreateTemp("", "config*.json")
if err != nil {
t.Fatal(err)
}
defer func() { _ = os.Remove(tmpfile.Name()) }()
content := `{
"bind": "127.0.0.1",
"port": 6380,
"max_clients": 500,
"read_timeout": 30,
"write_timeout": 60,
"max_message_bytes": 1048576,
"log_level": "warn"
}`
if _, err := tmpfile.WriteString(content); err != nil {
t.Fatal(err)
}
if err := tmpfile.Close(); err != nil {
t.Fatal(err)
}
cfg, err := Load(tmpfile.Name())
if err != nil {
t.Fatalf("Load() error: %v", err)
}
if cfg.Bind != "127.0.0.1" {
t.Errorf("Bind = %q, want 127.0.0.1", cfg.Bind)
}
if cfg.Port != 6380 {
t.Errorf("Port = %d, want 6380", cfg.Port)
}
if cfg.MaxClients != 500 {
t.Errorf("MaxClients = %d, want 500", cfg.MaxClients)
}
if cfg.ReadTimeout != 30 {
t.Errorf("ReadTimeout = %d, want 30", cfg.ReadTimeout)
}
if cfg.WriteTimeout != 60 {
t.Errorf("WriteTimeout = %d, want 60", cfg.WriteTimeout)
}
if cfg.MaxMsgBytes != 1048576 {
t.Errorf("MaxMsgBytes = %d, want 1048576", cfg.MaxMsgBytes)
}
if cfg.LogLevel != "warn" {
t.Errorf("LogLevel = %q, want warn", cfg.LogLevel)
}
}
func TestLoadInvalidFile(t *testing.T) {
_, err := Load("nonexistent.json")
if err == nil {
t.Error("Load() expected error for nonexistent file")
}
}
func TestLoadInvalidJSON(t *testing.T) {
tmpfile, err := os.CreateTemp("", "config*.json")
if err != nil {
t.Fatal(err)
}
defer func() { _ = os.Remove(tmpfile.Name()) }()
if _, err := tmpfile.WriteString("not valid json"); err != nil {
t.Fatal(err)
}
if err := tmpfile.Close(); err != nil {
t.Fatal(err)
}
_, err = Load(tmpfile.Name())
if err == nil {
t.Error("Load() expected error for invalid JSON")
}
}
func TestLoadAppliesDefaults(t *testing.T) {
tmpfile, err := os.CreateTemp("", "config*.json")
if err != nil {
t.Fatal(err)
}
defer func() { _ = os.Remove(tmpfile.Name()) }()
content := `{"bind": "127.0.0.1", "port": 6379}`
if _, err := tmpfile.WriteString(content); err != nil {
t.Fatal(err)
}
if err := tmpfile.Close(); err != nil {
t.Fatal(err)
}
cfg, err := Load(tmpfile.Name())
if err != nil {
t.Fatalf("Load() error: %v", err)
}
if cfg.MaxClients != DefaultMaxClients {
t.Errorf("MaxClients = %d, want default %d", cfg.MaxClients, DefaultMaxClients)
}
if cfg.LogLevel != DefaultLogLevel {
t.Errorf("LogLevel = %q, want default %q", cfg.LogLevel, DefaultLogLevel)
}
}
func TestAddr(t *testing.T) {
tests := []struct {
cfg *Config
want string
}{
// IPv4
{&Config{Bind: "0.0.0.0", Port: 6379}, "0.0.0.0:6379"},
{&Config{Bind: "127.0.0.1", Port: 6380}, "127.0.0.1:6380"},
{&Config{Bind: "192.168.1.100", Port: 8080}, "192.168.1.100:8080"},
// IPv6 — must be wrapped in brackets
{&Config{Bind: "::", Port: 6379}, "[::]:6379"},
{&Config{Bind: "::1", Port: 6379}, "[::1]:6379"},
{&Config{Bind: "2001:db8::1", Port: 8080}, "[2001:db8::1]:8080"},
{&Config{Bind: "fe80::1", Port: 6379}, "[fe80::1]:6379"},
}
for _, tt := range tests {
got := tt.cfg.Addr()
if got != tt.want {
t.Errorf("Addr(bind=%q, port=%d) = %q, want %q", tt.cfg.Bind, tt.cfg.Port, got, tt.want)
}
}
}
func TestReadTimeoutDuration(t *testing.T) {
tests := []struct {
name string
cfg *Config
expected string
}{
{"zero returns zero duration", &Config{ReadTimeout: 0}, "0s"},
{"30 seconds", &Config{ReadTimeout: 30}, "30s"},
{"120 seconds", &Config{ReadTimeout: 120}, "2m0s"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.cfg.ReadTimeoutDuration().String()
if got != tt.expected {
t.Errorf("ReadTimeoutDuration() = %s, want %s", got, tt.expected)
}
})
}
}
func TestWriteTimeoutDuration(t *testing.T) {
tests := []struct {
name string
cfg *Config
expected string
}{
{"zero returns zero duration", &Config{WriteTimeout: 0}, "0s"},
{"60 seconds", &Config{WriteTimeout: 60}, "1m0s"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.cfg.WriteTimeoutDuration().String()
if got != tt.expected {
t.Errorf("WriteTimeoutDuration() = %s, want %s", got, tt.expected)
}
})
}
}
func TestString(t *testing.T) {
cfg := &Config{
Bind: "0.0.0.0",
Port: 6379,
MaxClients: 10000,
ReadTimeout: 0,
WriteTimeout: 0,
MaxMsgBytes: 536870912,
LogLevel: "info",
}
str := cfg.String()
if len(str) == 0 {
t.Error("String() returned empty string")
}
}

View File

@@ -0,0 +1,24 @@
//go:build darwin || freebsd || openbsd || netbsd || dragonfly
// @tag config
// @tag platform:bsd
package config
import (
"os"
"path/filepath"
)
// DefaultConfigPaths returns the config file search paths for BSD and macOS.
// Looks up ~/.pekoe.json first, then /usr/local/etc/pekoe/config.json.
func DefaultConfigPaths() []string {
home, err := os.UserHomeDir()
if err != nil {
return []string{"/usr/local/etc/pekoe/config.json"}
}
return []string{
filepath.Join(home, ".pekoe.json"),
"/usr/local/etc/pekoe/config.json",
}
}

View File

@@ -0,0 +1,29 @@
//go:build linux
// @tag config
// @tag platform:linux
package config
import (
"os"
"path/filepath"
)
// DefaultConfigPaths returns the config file search paths for Linux.
// Root user: only /etc/pekoe/config.json
// Regular user: ~/.pekoe.json, then /etc/pekoe/config.json
func DefaultConfigPaths() []string {
if os.Getuid() == 0 {
return []string{"/etc/pekoe/config.json"}
}
home, err := os.UserHomeDir()
if err != nil {
return []string{"/etc/pekoe/config.json"}
}
return []string{
filepath.Join(home, ".pekoe.json"),
"/etc/pekoe/config.json",
}
}

View File

@@ -0,0 +1,11 @@
//go:build windows
// @tag config
// @tag platform:windows
package config
// DefaultConfigPaths returns no default paths for Windows.
// On Windows the --config flag is required.
func DefaultConfigPaths() []string {
return nil
}

170
internal/logger/logger.go Normal file
View File

@@ -0,0 +1,170 @@
// @tag logging
// @tag core
package logger
import (
"fmt"
"log"
"os"
"sync"
"time"
)
// Level represents the minimum severity for log messages.
type Level int
// Log severity levels from most verbose to most critical.
const (
LevelDebug Level = iota // Debug logs everything including diagnostic messages
LevelInfo // Info logs informational messages and above (default)
LevelWarn // Warn logs warnings, errors, and critical messages
LevelError // LevelError logs only errors and critical messages
)
var (
currentLevel = LevelInfo
logger = log.New(os.Stdout, "", 0)
logFile *os.File
logPath string
pid int
mu sync.Once
loggerMu sync.Mutex // Protects logger and logFile
)
// initPID gets the process ID for log messages
func initPID() {
pid = os.Getpid()
}
// formatLog formats a message with timestamp, log level, and pid.
// Format: timestamp level pekoe[pid]: message
func formatLog(level string, message string) string {
mu.Do(initPID)
// Timestamp format: "YYYY Mmm dd hh:mm:ss"
// e.g., "2024 Apr 2 14:30:15"
now := time.Now()
timestamp := now.Format("2006 Jan _2 15:04:05")
// Format: timestamp level pekoe[pid]: message
// e.g., "2024 Apr 2 14:30:15 INFO pekoe[12345]: event=startup"
return fmt.Sprintf("%s %s pekoe[%d]: %s", timestamp, level, pid, message)
}
// SetOutput configures the logger to write to the specified file.
// An empty path resets to stdout.
// Returns an error if the file cannot be opened/created.
// nolint: gosec
func SetOutput(path string) error {
loggerMu.Lock()
defer loggerMu.Unlock()
// Close existing log file if any
if logFile != nil {
_ = logFile.Close()
logFile = nil
}
if path == "" {
logPath = ""
logger = log.New(os.Stdout, "", 0)
return nil
}
file, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return err
}
logPath = path
logFile = file
logger = log.New(file, "", 0)
return nil
}
// Reopen reopens the log file.
// This is useful for log rotation - after the log file is moved/renamed,
// sending SIGUSR1 will cause the server to reopen the log file.
// Returns an error if the log file cannot be opened.
func Reopen() error {
loggerMu.Lock()
defer loggerMu.Unlock()
// If logging to stdout, nothing to reopen
if logPath == "" {
return nil
}
// Close existing file
if logFile != nil {
_ = logFile.Close()
logFile = nil
}
// Open the file again (create if doesn't exist)
//nolint: gosec // logPath is configured by admin, not user input
file, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
// Fall back to stdout if we can't reopen
logger = log.New(os.Stdout, "", 0)
return fmt.Errorf("failed to reopen log file %s: %w", logPath, err)
}
logFile = file
logger = log.New(file, "", 0)
return nil
}
// SetLevel configures the minimum log level from a string name.
// Accepted values: "debug", "info", "warn", "error".
func SetLevel(level string) {
switch level {
case "debug":
currentLevel = LevelDebug
case "info":
currentLevel = LevelInfo
case "warn":
currentLevel = LevelWarn
case "error":
currentLevel = LevelError
}
}
// Debug logs a message at debug level. Uses fmt.Sprintf formatting.
func Debug(format string, v ...any) {
if currentLevel <= LevelDebug {
msg := fmt.Sprintf(format, v...)
logger.Println(formatLog("DEBUG", msg))
}
}
// Info logs a message at info level. Uses fmt.Sprintf formatting.
func Info(format string, v ...any) {
if currentLevel <= LevelInfo {
msg := fmt.Sprintf(format, v...)
logger.Println(formatLog("INFO", msg))
}
}
// Warn logs a message at warning level. Uses fmt.Sprintf formatting.
func Warn(format string, v ...any) {
if currentLevel <= LevelWarn {
msg := fmt.Sprintf(format, v...)
logger.Println(formatLog("WARN", msg))
}
}
// Error logs a message at error level. Uses fmt.Sprintf formatting.
func Error(format string, v ...any) {
if currentLevel <= LevelError {
msg := fmt.Sprintf(format, v...)
logger.Println(formatLog("ERROR", msg))
}
}
// Fatal logs a message at critical level and terminates the process with exit code 1.
func Fatal(format string, v ...any) {
msg := fmt.Sprintf(format, v...)
logger.Println(formatLog("CRITICAL", msg))
os.Exit(1)
}

View File

@@ -0,0 +1,359 @@
package logger
import (
"log"
"os"
"path/filepath"
"strings"
"testing"
)
// resetLogger restores package-level state between tests.
func resetLogger() {
currentLevel = LevelInfo
loggerMu.Lock()
if logFile != nil {
_ = logFile.Close()
}
logFile = nil
logPath = ""
logger = log.New(os.Stdout, "", 0)
loggerMu.Unlock()
}
func TestSetLevel(t *testing.T) {
tests := []struct {
input string
want Level
}{
{"debug", LevelDebug},
{"info", LevelInfo},
{"warn", LevelWarn},
{"error", LevelError},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
resetLogger()
SetLevel(tt.input)
if currentLevel != tt.want {
t.Errorf("SetLevel(%q): got %d, want %d", tt.input, currentLevel, tt.want)
}
})
}
}
func TestSetLevelInvalid(t *testing.T) {
resetLogger()
SetLevel("info")
SetLevel("bogus")
if currentLevel != LevelInfo {
t.Errorf("invalid level changed current level: got %d, want %d", currentLevel, LevelInfo)
}
}
func TestSetLevelCaseSensitive(t *testing.T) {
resetLogger()
SetLevel("INFO") // uppercase should not match
if currentLevel != LevelInfo {
t.Errorf("uppercase level should not change level: got %d, want %d", currentLevel, LevelInfo)
}
}
func TestSetOutputToFile(t *testing.T) {
resetLogger()
dir := t.TempDir()
path := filepath.Join(dir, "test.log")
if err := SetOutput(path); err != nil {
t.Fatalf("SetOutput(%q): %v", path, err)
}
if logPath != path {
t.Errorf("logPath: got %q, want %q", logPath, path)
}
if logFile == nil {
t.Error("logFile should not be nil after SetOutput")
}
// Write something and verify it appears in the file
Info("test message")
if err := logFile.Sync(); err != nil {
t.Fatalf("sync: %v", err)
}
//nolint: gosec // test reads temp file in temp dir
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
if !strings.Contains(string(data), "test message") {
t.Errorf("log file should contain 'test message', got: %s", string(data))
}
// Close file before TempDir cleanup
resetLogger()
}
func TestSetOutputEmptyResetsToStdout(t *testing.T) {
resetLogger()
dir := t.TempDir()
path := filepath.Join(dir, "test.log")
if err := SetOutput(path); err != nil {
t.Fatalf("SetOutput: %v", err)
}
if err := SetOutput(""); err != nil {
t.Fatalf("SetOutput(\"\"): %v", err)
}
if logPath != "" {
t.Errorf("logPath should be empty, got %q", logPath)
}
if logFile != nil {
t.Error("logFile should be nil after reset to stdout")
}
}
func TestSetOutputInvalidPath(t *testing.T) {
resetLogger()
err := SetOutput("/nonexistent/dir/file.log")
if err == nil {
t.Error("SetOutput with invalid path should return error")
}
}
func TestReopenNoFile(t *testing.T) {
resetLogger()
// Reopen when logging to stdout should be a no-op
if err := Reopen(); err != nil {
t.Errorf("Reopen with no file should not error: %v", err)
}
}
func TestReopen(t *testing.T) {
resetLogger()
dir := t.TempDir()
path := filepath.Join(dir, "test.log")
if err := SetOutput(path); err != nil {
t.Fatalf("SetOutput: %v", err)
}
// Write initial message
Info("before reopen")
_ = logFile.Sync()
// Close and remove the file to simulate log rotation
_ = logFile.Close()
_ = os.Remove(path)
// Reopen should recreate the file
if err := Reopen(); err != nil {
t.Fatalf("Reopen: %v", err)
}
if logFile == nil {
t.Error("logFile should not be nil after reopen")
}
// Write after reopen
Info("after reopen")
_ = logFile.Sync()
//nolint: gosec // test reads temp file in temp dir
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
content := string(data)
if !strings.Contains(content, "after reopen") {
t.Errorf("log should contain 'after reopen', got: %s", content)
}
if strings.Contains(content, "before reopen") {
t.Error("log should not contain 'before reopen' after file was removed")
}
// Close file before TempDir cleanup
resetLogger()
}
func TestFormatLog(t *testing.T) {
// Force PID init
mu.Do(initPID)
msg := formatLog("INFO", "test key=value")
// Verify format: timestamp level pekoe[pid]: message
if !strings.Contains(msg, "INFO pekoe[") {
t.Errorf("expected 'INFO pekoe[', got: %s", msg)
}
if !strings.Contains(msg, "]: test key=value") {
t.Errorf("expected message content, got: %s", msg)
}
}
func TestFormatLogLevels(t *testing.T) {
mu.Do(initPID)
levels := []string{"DEBUG", "INFO", "WARN", "ERROR", "CRITICAL"}
for _, level := range levels {
msg := formatLog(level, "x")
if !strings.Contains(msg, level+" pekoe[") {
t.Errorf("level %s: expected '%s pekoe[', got: %s", level, level, msg)
}
}
}
func TestLevelFiltering(t *testing.T) {
tests := []struct {
level string
logFunc func(string, ...interface{})
name string
want bool // whether message should appear
}{
// At debug level, all messages should appear
{"debug", Debug, "debug_at_debug", true},
{"debug", Info, "info_at_debug", true},
{"debug", Warn, "warn_at_debug", true},
{"debug", Error, "error_at_debug", true},
// At info level, debug should be filtered
{"info", Debug, "debug_at_info", false},
{"info", Info, "info_at_info", true},
{"info", Warn, "warn_at_info", true},
{"info", Error, "error_at_info", true},
// At warn level, debug and info should be filtered
{"warn", Debug, "debug_at_warn", false},
{"warn", Info, "info_at_warn", false},
{"warn", Warn, "warn_at_warn", true},
{"warn", Error, "error_at_warn", true},
// At error level, only error should appear
{"error", Debug, "debug_at_error", false},
{"error", Info, "info_at_error", false},
{"error", Warn, "warn_at_error", false},
{"error", Error, "error_at_error", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resetLogger()
dir := t.TempDir()
path := filepath.Join(dir, "filter.log")
if err := SetOutput(path); err != nil {
t.Fatalf("SetOutput: %v", err)
}
SetLevel(tt.level)
tt.logFunc(tt.name)
_ = logFile.Sync()
//nolint: gosec // test reads temp file in temp dir
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
content := string(data)
contains := strings.Contains(content, tt.name)
if tt.want && !contains {
t.Errorf("level=%s: expected %s to appear in log, got: %s", tt.level, tt.name, content)
}
if !tt.want && contains {
t.Errorf("level=%s: expected %s to be filtered, but it appeared in log", tt.level, tt.name)
}
// Close file before TempDir cleanup
resetLogger()
})
}
}
func TestLogFormatIncludesKeyValues(t *testing.T) {
resetLogger()
dir := t.TempDir()
path := filepath.Join(dir, "format.log")
if err := SetOutput(path); err != nil {
t.Fatalf("SetOutput: %v", err)
}
SetLevel("debug")
Info("event=test key=value count=%d", 42)
_ = logFile.Sync()
//nolint: gosec // test reads temp file in temp dir
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
content := string(data)
// Should contain the formatted message
if !strings.Contains(content, "event=test") {
t.Errorf("expected 'event=test' in log, got: %s", content)
}
if !strings.Contains(content, "key=value") {
t.Errorf("expected 'key=value' in log, got: %s", content)
}
if !strings.Contains(content, "count=42") {
t.Errorf("expected 'count=42' in log, got: %s", content)
}
// Close file before TempDir cleanup
resetLogger()
}
func TestMultipleSetOutputSwitches(t *testing.T) {
resetLogger()
dir := t.TempDir()
path1 := filepath.Join(dir, "first.log")
path2 := filepath.Join(dir, "second.log")
// Write to first file
if err := SetOutput(path1); err != nil {
t.Fatalf("SetOutput path1: %v", err)
}
SetLevel("info")
Info("first file")
_ = logFile.Sync()
// Switch to second file (SetOutput closes the first file)
if err := SetOutput(path2); err != nil {
t.Fatalf("SetOutput path2: %v", err)
}
Info("second file")
_ = logFile.Sync()
// Verify first file has only its message
//nolint: gosec // test reads temp files in temp dir
data1, err := os.ReadFile(path1)
if err != nil {
t.Fatalf("ReadFile path1: %v", err)
}
if !strings.Contains(string(data1), "first file") {
t.Error("first log should contain 'first file'")
}
if strings.Contains(string(data1), "second file") {
t.Error("first log should not contain 'second file'")
}
// Verify second file has only its message
//nolint: gosec // test reads temp files in temp dir
data2, err := os.ReadFile(path2)
if err != nil {
t.Fatalf("ReadFile path2: %v", err)
}
if !strings.Contains(string(data2), "second file") {
t.Error("second log should contain 'second file'")
}
if strings.Contains(string(data2), "first file") {
t.Error("second log should not contain 'first file'")
}
// Close file before TempDir cleanup
resetLogger()
}

View File

@@ -0,0 +1,252 @@
// @tag protocol
// @tag core
package protocol
import (
"bufio"
"bytes"
"errors"
"fmt"
"io"
"strconv"
)
const (
CRLF = "\r\n"
MaxBulkLen = 512 * 1024 * 1024 // 512 MB limit for bulk strings - prevents memory exhaustion
MaxArrayLen = 1024 * 1024 // 1M elements max - prevents array parsing attacks
)
// Protocol errors - separated for clear error handling
var (
ErrProtocolError = errors.New("protocol error")
ErrNotEnoughData = errors.New("not enough data")
)
// Value represents any RESP value - using interface for simplicity with type switches.
// This mirrors Redis protocol flexibility while keeping implementation simple.
type Value any
// Error represents an error response from Redis (-ERR...)
type Error struct {
Message string
}
func (e Error) Error() string {
return e.Message
}
// Array represents an array of values (*...) - used for commands and responses.
type Array []Value
// Bulk represents a bulk string ($...) - binary safe, used for most arguments.
type Bulk []byte
func (b Bulk) String() string {
return string(b)
}
// Int represents an integer response (:12345) - used for counts and numeric responses.
type Int int64
// Status represents a simple string response (+OK) - used for simple confirmations.
type Status string
func (s Status) String() string {
return string(s)
}
// Parse reads and parses a single RESP value from the reader.
// This is the main entry point for parsing incoming Redis commands.
// Uses bufio.Reader for efficient I/O with built-in buffering.
func Parse(reader *bufio.Reader) (Value, error) {
// Read first line to determine message type
line, err := readLine(reader)
if err != nil {
return nil, err
}
if len(line) == 0 {
return nil, ErrProtocolError
}
// Route based on first character - this is the RESP type indicator
switch line[0] {
case '+': // Simple string: +OK\r\n
return Status(line[1:]), nil
case '-': // Error: -ERR message\r\n
return Error{Message: string(line[1:])}, nil
case ':': // Integer: :12345\r\n
n, err := strconv.ParseInt(string(line[1:]), 10, 64)
if err != nil {
return nil, ErrProtocolError
}
return Int(n), nil
case '$': // Bulk string: $6\r\nfoobar\r\n or $-1 for null
return parseBulk(reader, line)
case '*': // Array: *2\r\n$3\r\nfoo\r\n$3\r\nbar\r\n
return parseArray(reader, line)
default:
return nil, fmt.Errorf("unknown type: %c", line[0])
}
}
// readLine reads a single line from the reader, handling CRLF.
// Uses ReadLine for efficiency - it handles buffering internally.
func readLine(reader *bufio.Reader) ([]byte, error) {
line, isPrefix, err := reader.ReadLine()
if err != nil {
return nil, err
}
if isPrefix {
return nil, fmt.Errorf("line too long")
}
return line, nil
}
// parseBulk handles bulk string parsing ($<length>\r\n<data>\r\n).
// Bulk strings are used for binary-safe string arguments and responses.
func parseBulk(reader *bufio.Reader, line []byte) (Value, error) {
// Parse the length field after the $ prefix
size, err := strconv.Atoi(string(line[1:]))
if err != nil {
return nil, ErrProtocolError
}
// $-1 means null bulk string - legitimate empty value
if size == -1 {
return nil, nil
}
// Validate size to prevent memory exhaustion attacks
if size < 0 || size > MaxBulkLen {
return nil, fmt.Errorf("bulk string too large: %d bytes", size)
}
// Read exactly size bytes plus CRLF delimiter
// Using ReadFull ensures we get all data or an error
buf := make([]byte, size+len(CRLF))
_, err = io.ReadFull(reader, buf)
if err != nil {
return nil, err
}
// Verify proper CRLF termination - prevents protocol injection
if !bytes.HasSuffix(buf, []byte(CRLF)) {
return nil, ErrProtocolError
}
return Bulk(buf[:size]), nil
}
// parseArray handles array parsing (*<count>\r\n<elements>...).
// Arrays are the fundamental structure for Redis commands.
func parseArray(reader *bufio.Reader, line []byte) (Value, error) {
count, err := strconv.Atoi(string(line[1:]))
if err != nil {
return nil, ErrProtocolError
}
// *-1 means null array - rare but valid
if count == -1 {
return nil, nil
}
// Limit array size to prevent DoS via large arrays
if count < 0 || count > MaxArrayLen {
return nil, fmt.Errorf("array too large: %d elements", count)
}
// Pre-allocate with expected size to reduce reallocations
result := make([]Value, count)
for i := range count {
val, err := Parse(reader)
if err != nil {
return nil, err
}
result[i] = val
}
return Array(result), nil
}
// Response builders - these create RESP-encoded responses to send to clients
// MakeError creates an error response (-ERR message).
func MakeError(msg string) []byte {
return []byte("-" + msg + CRLF)
}
// MakeStatus creates a simple status response (+OK).
func MakeStatus(status string) []byte {
return []byte("+" + status + CRLF)
}
// MakeInt creates an integer response (:12345).
func MakeInt(n int64) []byte {
return []byte(":" + strconv.FormatInt(n, 10) + CRLF)
}
// MakeBulk creates a bulk string response ($<len>\r\n<data>\r\n).
// Handles nil gracefully by returning null bulk string.
func MakeBulk(data []byte) []byte {
if data == nil {
return []byte("$-1" + CRLF)
}
return []byte("$" + strconv.Itoa(len(data)) + CRLF + string(data) + CRLF)
}
// MakeMultiBulk creates a RESP array from Value slices - used for pub/sub messages.
func MakeMultiBulk(values []Value) []byte {
var buf bytes.Buffer
buf.WriteString("*" + strconv.Itoa(len(values)) + CRLF)
for _, v := range values {
switch val := v.(type) {
case nil:
buf.WriteString("$-1\r\n")
case string:
buf.Write(MakeBulk([]byte(val)))
case []byte:
buf.Write(MakeBulk(val))
case int64:
buf.Write(MakeInt(val))
case Int:
buf.Write(MakeInt(int64(val)))
case Error:
buf.Write(MakeError(val.Message))
case Status:
buf.Write(MakeStatus(string(val)))
default:
buf.Write(MakeBulk([]byte(fmt.Sprintf("%v", val))))
}
}
return buf.Bytes()
}
// ToBytes converts a parsed Value back to its RESP wire format.
// This is useful for converting parsed commands back to bytes for relay or logging.
// nolint: cyclop
func ToBytes(v Value) []byte {
if v == nil {
return MakeBulk(nil)
}
switch val := v.(type) {
case string:
return MakeBulk([]byte(val))
case []byte:
return MakeBulk(val)
case Bulk:
return MakeBulk([]byte(val))
case int64:
return MakeInt(val)
case Int:
return MakeInt(int64(val))
case Error:
return MakeError(val.Message)
case Status:
return MakeStatus(string(val))
case Array:
return MakeMultiBulk(val)
default:
return MakeBulk([]byte(fmt.Sprintf("%v", val)))
}
}

View File

@@ -0,0 +1,206 @@
package protocol
import (
"bufio"
"bytes"
"testing"
)
func TestParseSimpleString(t *testing.T) {
tests := []struct {
input string
expected Status
}{
{"+OK\r\n", "OK"},
{"+PONG\r\n", "PONG"},
{"+\r\n", ""},
}
for _, tt := range tests {
reader := bufio.NewReader(bytes.NewReader([]byte(tt.input)))
val, err := Parse(reader)
if err != nil {
t.Errorf("Parse(%q) error: %v", tt.input, err)
continue
}
status, ok := val.(Status)
if !ok {
t.Errorf("Parse(%q) got %T, want Status", tt.input, val)
continue
}
if status != tt.expected {
t.Errorf("Parse(%q) = %q, want %q", tt.input, status, tt.expected)
}
}
}
func TestParseError(t *testing.T) {
input := "-ERR unknown command\r\n"
reader := bufio.NewReader(bytes.NewReader([]byte(input)))
val, err := Parse(reader)
if err != nil {
t.Fatalf("Parse error: %v", err)
}
errVal, ok := val.(Error)
if !ok {
t.Fatalf("got %T, want Error", val)
}
if errVal.Message != "ERR unknown command" {
t.Errorf("got %q, want 'ERR unknown command'", errVal.Message)
}
}
func TestParseInteger(t *testing.T) {
tests := []struct {
input string
expected Int
}{
{":0\r\n", 0},
{":12345\r\n", 12345},
{":-100\r\n", -100},
}
for _, tt := range tests {
reader := bufio.NewReader(bytes.NewReader([]byte(tt.input)))
val, err := Parse(reader)
if err != nil {
t.Errorf("Parse(%q) error: %v", tt.input, err)
continue
}
n, ok := val.(Int)
if !ok {
t.Errorf("Parse(%q) got %T, want Int", tt.input, val)
continue
}
if n != tt.expected {
t.Errorf("Parse(%q) = %d, want %d", tt.input, n, tt.expected)
}
}
}
func TestParseBulkString(t *testing.T) {
tests := []struct {
input string
expected []byte
}{
{"$5\r\nhello\r\n", []byte("hello")},
{"$0\r\n\r\n", []byte("")},
{"$-1\r\n", nil},
}
for _, tt := range tests {
reader := bufio.NewReader(bytes.NewReader([]byte(tt.input)))
val, err := Parse(reader)
if err != nil {
t.Errorf("Parse(%q) error: %v", tt.input, err)
continue
}
if tt.expected == nil {
if val != nil {
t.Errorf("Parse(%q) = %v, want nil", tt.input, val)
}
continue
}
bulk, ok := val.(Bulk)
if !ok {
t.Errorf("Parse(%q) got %T, want Bulk", tt.input, val)
continue
}
if string(bulk) != string(tt.expected) {
t.Errorf("Parse(%q) = %q, want %q", tt.input, string(bulk), string(tt.expected))
}
}
}
func TestParseArray(t *testing.T) {
input := "*2\r\n$3\r\nfoo\r\n$3\r\nbar\r\n"
reader := bufio.NewReader(bytes.NewReader([]byte(input)))
val, err := Parse(reader)
if err != nil {
t.Fatalf("Parse error: %v", err)
}
arr, ok := val.(Array)
if !ok {
t.Fatalf("got %T, want Array", val)
}
if len(arr) != 2 {
t.Errorf("array length = %d, want 2", len(arr))
}
}
func TestParseEmptyArray(t *testing.T) {
input := "*0\r\n"
reader := bufio.NewReader(bytes.NewReader([]byte(input)))
val, err := Parse(reader)
if err != nil {
t.Fatalf("Parse error: %v", err)
}
arr, ok := val.(Array)
if !ok {
t.Fatalf("got %T, want Array", val)
}
if len(arr) != 0 {
t.Errorf("array length = %d, want 0", len(arr))
}
}
func TestMakeStatus(t *testing.T) {
result := MakeStatus("OK")
expected := "+OK\r\n"
if string(result) != expected {
t.Errorf("MakeStatus() = %q, want %q", string(result), expected)
}
}
func TestMakeInt(t *testing.T) {
result := MakeInt(42)
expected := ":42\r\n"
if string(result) != expected {
t.Errorf("MakeInt() = %q, want %q", string(result), expected)
}
}
func TestMakeBulk(t *testing.T) {
tests := []struct {
input []byte
expected string
}{
{[]byte("hello"), "$5\r\nhello\r\n"},
{[]byte(""), "$0\r\n\r\n"},
{nil, "$-1\r\n"},
}
for _, tt := range tests {
result := MakeBulk(tt.input)
if string(result) != tt.expected {
t.Errorf("MakeBulk(%q) = %q, want %q", tt.input, string(result), tt.expected)
}
}
}
func TestMakeMultiBulk(t *testing.T) {
result := MakeMultiBulk([]Value{Bulk("message"), Bulk("mychannel"), Bulk("hello")})
if len(result) == 0 {
t.Error("MakeMultiBulk returned empty")
}
}
func TestToBytes(t *testing.T) {
tests := []struct {
input Value
expected string
}{
{nil, "$-1\r\n"},
{Status("OK"), "+OK\r\n"},
{Int(42), ":42\r\n"},
{Error{Message: "ERR"}, "-ERR\r\n"},
{Bulk("test"), "$4\r\ntest\r\n"},
}
for _, tt := range tests {
result := ToBytes(tt.input)
if string(result) != tt.expected {
t.Errorf("ToBytes(%v) = %q, want %q", tt.input, string(result), tt.expected)
}
}
}

257
internal/pubsub/client.go Normal file
View File

@@ -0,0 +1,257 @@
// @tag pubsub
// @tag core
package pubsub
import (
"fmt"
"net"
"sync"
"sync/atomic"
"time"
"github.com/eleksir/pekoe/internal/protocol"
)
// ReplyMode controls client reply behavior
type ReplyMode int
const (
ReplyOn ReplyMode = iota // Normal reply mode - send all responses
ReplyOff // Reply off - suppress all responses
ReplySkip // Reply skip - skip the next response only
)
// Global client ID counter - each client gets a unique ID
var clientIDCounter atomic.Uint64
// Client represents a subscriber connection to the pubsub system.
// It maintains the connection (net.Conn) and tracks which channels/patterns the client is subscribed to.
type Client struct {
id uint64 // Unique client ID
conn net.Conn
name string // Connection name set via CLIENT SETNAME
replyMode ReplyMode // Reply mode (ON/OFF/SKIP)
flags uint64 // Client flags (e.g., paused, no-touch)
created time.Time // When the client connected
lastActive time.Time // Last activity time
subscribed map[string]struct{}
patterns map[string]struct{}
mu sync.RWMutex
writeMu sync.Mutex // Protects concurrent writes to conn
}
// NewClient creates a new pubsub client wrapping a connection.
func NewClient(conn net.Conn) *Client {
now := time.Now()
return &Client{
id: clientIDCounter.Add(1),
conn: conn,
name: "",
replyMode: ReplyOn,
flags: 0,
created: now,
lastActive: now,
subscribed: make(map[string]struct{}),
patterns: make(map[string]struct{}),
}
}
// ID returns the unique client ID
func (c *Client) ID() uint64 {
return c.id
}
// SetName sets the connection name
func (c *Client) SetName(name string) {
c.mu.Lock()
defer c.mu.Unlock()
c.name = name
}
// GetName returns the connection name
func (c *Client) GetName() string {
c.mu.RLock()
defer c.mu.RUnlock()
return c.name
}
// SetReplyMode sets the reply mode
func (c *Client) SetReplyMode(mode ReplyMode) {
c.mu.Lock()
defer c.mu.Unlock()
c.replyMode = mode
}
// GetReplyMode returns the current reply mode
func (c *Client) GetReplyMode() ReplyMode {
c.mu.RLock()
defer c.mu.RUnlock()
return c.replyMode
}
// UpdateActivity updates the last activity timestamp
func (c *Client) UpdateActivity() {
c.mu.Lock()
defer c.mu.Unlock()
c.lastActive = time.Now()
}
// GetChannels returns a copy of the channels this client is subscribed to.
// Returns a slice to prevent external modification.
func (c *Client) GetChannels() []string {
c.mu.RLock()
defer c.mu.RUnlock()
channels := make([]string, 0, len(c.subscribed))
for ch := range c.subscribed {
channels = append(channels, ch)
}
return channels
}
// GetPatterns returns a copy of the patterns this client is subscribed to.
func (c *Client) GetPatterns() []string {
c.mu.RLock()
defer c.mu.RUnlock()
patterns := make([]string, 0, len(c.patterns))
for p := range c.patterns {
patterns = append(patterns, p)
}
return patterns
}
// SubsCount returns the number of channels this client is subscribed to.
func (c *Client) SubsCount() int {
c.mu.RLock()
defer c.mu.RUnlock()
return len(c.subscribed)
}
// PatternSubsCount returns the number of patterns this client is subscribed to.
func (c *Client) PatternSubsCount() int {
c.mu.RLock()
defer c.mu.RUnlock()
return len(c.patterns)
}
// AddSubscription tracks that this client has subscribed to a channel.
func (c *Client) AddSubscription(channel string) {
c.mu.Lock()
defer c.mu.Unlock()
c.subscribed[channel] = struct{}{}
c.lastActive = time.Now()
}
// RemoveSubscription removes a channel from the client's subscriptions.
func (c *Client) RemoveSubscription(channel string) {
c.mu.Lock()
defer c.mu.Unlock()
delete(c.subscribed, channel)
c.lastActive = time.Now()
}
// AddPatternSubscription tracks that this client has subscribed to a pattern.
func (c *Client) AddPatternSubscription(pattern string) {
c.mu.Lock()
defer c.mu.Unlock()
c.patterns[pattern] = struct{}{}
c.lastActive = time.Now()
}
// RemovePatternSubscription removes a pattern from the client's subscriptions.
func (c *Client) RemovePatternSubscription(pattern string) {
c.mu.Lock()
defer c.mu.Unlock()
delete(c.patterns, pattern)
c.lastActive = time.Now()
}
// SendMessage sends a pubsub message to the client in RESP format.
// Format: ["message", channel, payload]
// Returns error if the write fails - caller will handle cleanup.
func (c *Client) sendMessage(channel string, message []byte) error {
response := protocol.MakeMultiBulk([]protocol.Value{
protocol.Bulk("message"),
protocol.Bulk(channel),
protocol.Bulk(message),
})
c.writeMu.Lock()
_, err := c.conn.Write(response)
c.writeMu.Unlock()
c.UpdateActivity()
return err
}
// SendPMessage sends a pubsub message matching a pattern.
// Format: ["pmessage", pattern, channel, payload]
func (c *Client) sendPMessage(pattern string, channel string, message []byte) error {
response := protocol.MakeMultiBulk([]protocol.Value{
protocol.Bulk("pmessage"),
protocol.Bulk(pattern),
protocol.Bulk(channel),
protocol.Bulk(message),
})
c.writeMu.Lock()
_, err := c.conn.Write(response)
c.writeMu.Unlock()
c.UpdateActivity()
return err
}
// RemoteAddr returns the remote address of the client for logging.
func (c *Client) RemoteAddr() string {
return c.conn.RemoteAddr().String()
}
// LocalAddr returns the local address of the client connection.
func (c *Client) LocalAddr() string {
return c.conn.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=
func (c *Client) InfoString() string {
c.mu.RLock()
defer c.mu.RUnlock()
age := int(time.Since(c.created).Seconds())
idle := int(time.Since(c.lastActive).Seconds())
// Determine flags
flags := ""
if c.replyMode == ReplyOff {
flags += "O"
}
if c.replyMode == ReplySkip {
flags += "S"
}
// 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",
c.id,
c.RemoteAddr(),
c.LocalAddr(),
age,
idle,
flags,
len(c.subscribed),
len(c.patterns),
)
}
// ListString returns CLIENT LIST formatted string
// Format: id= addr= ... (same as INFO but one line per client)
func (c *Client) ListString() string {
return c.InfoString()
}
// Close closes the client's connection.
// Called when the client disconnects or on error.
func (c *Client) Close() error {
c.writeMu.Lock()
defer c.writeMu.Unlock()
return c.conn.Close()
}

463
internal/pubsub/hub.go Normal file
View File

@@ -0,0 +1,463 @@
// @tag pubsub
// @tag core
package pubsub
import (
"fmt"
"regexp"
"slices"
"strings"
"sync"
)
// Hub maintains the mapping between channels and their subscribers.
type Hub struct {
mu sync.RWMutex
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
}
// patternEntry holds compiled regex and subscribers for a pattern.
type patternEntry struct {
regex *regexp.Regexp
clients map[*Client]struct{}
}
// patternCache caches compiled regex patterns to avoid recompilation.
type patternCacheType struct {
sync.Map
}
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),
}
}
// RegisterClient adds a client to the registry
func (h *Hub) RegisterClient(client *Client) {
h.mu.Lock()
defer h.mu.Unlock()
h.clients[client.ID()] = client
}
// UnregisterClient removes a client from the registry
func (h *Hub) UnregisterClient(client *Client) {
h.mu.Lock()
defer h.mu.Unlock()
delete(h.clients, client.ID())
}
// GetClient returns a client by ID, or nil if not found
func (h *Hub) GetClient(id uint64) *Client {
h.mu.RLock()
defer h.mu.RUnlock()
return h.clients[id]
}
// GetAllClients returns a slice of all registered clients
func (h *Hub) GetAllClients() []*Client {
h.mu.RLock()
defer h.mu.RUnlock()
clients := make([]*Client, 0, len(h.clients))
for _, c := range h.clients {
clients = append(clients, c)
}
return clients
}
// ClientCount returns the number of connected clients
func (h *Hub) ClientCount() int {
h.mu.RLock()
defer h.mu.RUnlock()
return len(h.clients)
}
// Subscribe adds a client to a channel.
func (h *Hub) Subscribe(client *Client, channel string) bool {
if channel == "" {
return false
}
// Must hold hub mutex when modifying the global subs map
h.mu.Lock()
defer h.mu.Unlock()
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 {
return false
}
clients[client] = struct{}{}
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 == "" {
return false
}
// Must hold hub mutex when modifying the global subs map
h.mu.Lock()
defer h.mu.Unlock()
clients, exists := h.subs[channel]
if !exists {
return false
}
if _, ok := clients[client]; !ok {
return false
}
delete(clients, client)
if len(clients) == 0 {
delete(h.subs, channel)
}
client.RemoveSubscription(channel)
return true
}
// UnsubscribeAll removes a client from all subscribed channels and patterns.
func (h *Hub) UnsubscribeAll(client *Client) {
channels := client.GetChannels()
for _, channel := range channels {
h.Unsubscribe(client, channel)
}
patterns := client.GetPatterns()
for _, pattern := range patterns {
h.PatternUnsubscribe(client, pattern)
}
}
// Publish sends a message to all subscribers of a channel.
func (h *Hub) Publish(channel string, message []byte) int {
// Copy clients under lock, then release lock before sending
h.mu.RLock()
clients, exists := h.subs[channel]
if !exists || len(clients) == 0 {
h.mu.RUnlock()
return 0
}
clientList := make([]*Client, 0, len(clients))
for client := range clients {
clientList = append(clientList, client)
}
h.mu.RUnlock()
count := 0
for _, client := range clientList {
if err := client.sendMessage(channel, message); err != nil {
h.removeClientFromChannel(client, channel)
continue
}
count++
}
return count
}
// PublishToPattern sends a message to all clients subscribed to patterns matching the channel.
// nolint: cyclop
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
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)
}
}
}
if len(matchedPatterns) == 0 {
h.mu.RUnlock()
return 0
}
type match struct {
client *Client
pattern string
}
var matches []match
for _, pattern := range matchedPatterns {
if entry, ok := h.patterns[pattern]; ok {
for client := range entry.clients {
matches = append(matches, match{client: client, pattern: pattern})
}
}
}
h.mu.RUnlock()
count := 0
for _, m := range matches {
if err := m.client.sendPMessage(m.pattern, channel, message); err != nil {
h.removeClientFromPattern(m.client, m.pattern)
continue
}
count++
}
return count
}
// PatternSubscribe adds a client to a pattern subscription.
func (h *Hub) PatternSubscribe(client *Client, pattern string) (bool, error) {
regex, err := getOrCompilePattern(pattern)
if err != nil {
return false, fmt.Errorf("invalid pattern: %v", err)
}
h.mu.Lock()
defer h.mu.Unlock()
entry, exists := h.patterns[pattern]
if !exists {
entry = &patternEntry{
regex: regex,
clients: make(map[*Client]struct{}),
}
h.patterns[pattern] = entry
}
if _, already := entry.clients[client]; already {
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)
}
}
client.AddPatternSubscription(pattern)
return true, nil
}
// getOrCompilePattern returns a cached compiled regex or compiles and caches a new one.
func getOrCompilePattern(pattern string) (*regexp.Regexp, error) {
if cached, ok := patternCache.Load(pattern); ok {
return cached.(*regexp.Regexp), nil
}
regex, err := compilePattern(pattern)
if err != nil {
return nil, err
}
patternCache.Store(pattern, regex)
return regex, nil
}
func compilePattern(pattern string) (*regexp.Regexp, error) {
var buf strings.Builder
for i := 0; i < len(pattern); i++ {
switch pattern[i] {
case '*':
buf.WriteString(".*")
case '?':
buf.WriteByte('.')
case '\\':
if i+1 < len(pattern) {
buf.WriteByte('\\')
buf.WriteByte(pattern[i+1])
i++
}
case '.', '+', '(', ')', '[', ']', '{', '}', '^', '$', '|':
buf.WriteByte('\\')
buf.WriteByte(pattern[i])
default:
buf.WriteByte(pattern[i])
}
}
return regexp.Compile("^" + buf.String() + "$")
}
// PatternUnsubscribe removes a client from a pattern subscription.
func (h *Hub) PatternUnsubscribe(client *Client, pattern string) bool {
h.mu.Lock()
defer h.mu.Unlock()
entry, exists := h.patterns[pattern]
if !exists {
return false
}
if _, ok := entry.clients[client]; !ok {
return false
}
delete(entry.clients, client)
if len(entry.clients) == 0 {
delete(h.patterns, pattern)
// Remove pattern from index
h.removePatternFromIndex(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()
defer h.mu.RUnlock()
count := 0
for _, entry := range h.patterns {
count += len(entry.clients)
}
return count
}
// RemoveClientFromChannel handles cleanup when a client's connection fails.
func (h *Hub) removeClientFromChannel(client *Client, channel string) {
h.mu.Lock()
defer h.mu.Unlock()
if clients, ok := h.subs[channel]; ok {
if _, exists := clients[client]; exists {
delete(clients, client)
if len(clients) == 0 {
delete(h.subs, channel)
}
}
}
client.RemoveSubscription(channel)
}
// RemoveClientFromPattern handles cleanup when a client's connection fails for patterns.
func (h *Hub) removeClientFromPattern(client *Client, pattern string) {
h.mu.Lock()
defer h.mu.Unlock()
if entry, ok := h.patterns[pattern]; ok {
if _, exists := entry.clients[client]; exists {
delete(entry.clients, client)
if len(entry.clients) == 0 {
delete(h.patterns, pattern)
}
}
}
client.RemovePatternSubscription(pattern)
}
// ChannelCount returns the number of subscribers to a specific channel.
func (h *Hub) ChannelCount(channel string) int {
h.mu.RLock()
defer h.mu.RUnlock()
if clients, ok := h.subs[channel]; ok {
return len(clients)
}
return 0
}
// TotalChannels returns the total number of active channels.
func (h *Hub) TotalChannels() int {
h.mu.RLock()
defer h.mu.RUnlock()
return len(h.subs)
}
// ChannelList returns a list of all active channel names, optionally filtered by pattern.
func (h *Hub) ChannelList(pattern string) []string {
h.mu.RLock()
defer h.mu.RUnlock()
if pattern == "" {
channels := make([]string, 0, len(h.subs))
for ch := range h.subs {
channels = append(channels, ch)
}
return channels
}
// Compile pattern once outside the loop (using cache)
regex, err := getOrCompilePattern(pattern)
if err != nil {
return nil
}
channels := make([]string, 0, len(h.subs))
for ch := range h.subs {
if regex.MatchString(ch) {
channels = append(channels, ch)
}
}
return channels
}

View File

@@ -0,0 +1,333 @@
package pubsub
import (
"net"
"testing"
"time"
)
type mockConn struct {
data []byte
pos int
}
func (m *mockConn) Read(b []byte) (n int, err error) {
if m.pos >= len(m.data) {
return 0, nil
}
n = copy(b, m.data[m.pos:])
m.pos += n
return n, nil
}
func (m *mockConn) Write(b []byte) (n int, err error) {
m.data = append(m.data, b...)
return len(b), nil
}
func (m *mockConn) Close() error { return nil }
func (m *mockConn) LocalAddr() net.Addr { return nil }
func (m *mockConn) RemoteAddr() net.Addr { return nil }
func (m *mockConn) SetDeadline(time.Time) error { return nil }
func (m *mockConn) SetReadDeadline(time.Time) error { return nil }
func (m *mockConn) SetWriteDeadline(time.Time) error { return nil }
func NewMockClient() *Client {
return NewClient(&mockConn{data: []byte{}})
}
func TestNewHub(t *testing.T) {
hub := NewHub()
if hub == nil {
t.Fatal("NewHub() returned nil")
}
if hub.subs == nil {
t.Error("hub.subs is nil")
}
}
func TestHubSubscribe(t *testing.T) {
hub := NewHub()
client := NewMockClient()
// Subscribe to a channel
ok := hub.Subscribe(client, "news")
if !ok {
t.Error("Subscribe() returned false for new subscription")
}
// Check subscriber count
count := hub.ChannelCount("news")
if count != 1 {
t.Errorf("ChannelCount() = %d, want 1", count)
}
// Duplicate subscription should return false
ok = hub.Subscribe(client, "news")
if ok {
t.Error("Subscribe() should return false for duplicate subscription")
}
}
func TestHubUnsubscribe(t *testing.T) {
hub := NewHub()
client := NewMockClient()
hub.Subscribe(client, "news")
ok := hub.Unsubscribe(client, "news")
if !ok {
t.Error("Unsubscribe() returned false")
}
count := hub.ChannelCount("news")
if count != 0 {
t.Errorf("ChannelCount() = %d, want 0", count)
}
}
func TestHubUnsubscribeNotSubscribed(t *testing.T) {
hub := NewHub()
client := NewMockClient()
ok := hub.Unsubscribe(client, "nonexistent")
if ok {
t.Error("Unsubscribe() should return false for non-subscribed channel")
}
}
func TestHubUnsubscribeAll(t *testing.T) {
hub := NewHub()
client := NewMockClient()
hub.Subscribe(client, "news")
hub.Subscribe(client, "sports")
hub.Subscribe(client, "weather")
hub.UnsubscribeAll(client)
if hub.ChannelCount("news") != 0 {
t.Error("news channel should be empty")
}
if hub.ChannelCount("sports") != 0 {
t.Error("sports channel should be empty")
}
if hub.ChannelCount("weather") != 0 {
t.Error("weather channel should be empty")
}
}
func TestHubPublish(t *testing.T) {
hub := NewHub()
client1 := NewMockClient()
client2 := NewMockClient()
hub.Subscribe(client1, "news")
hub.Subscribe(client2, "news")
count := hub.Publish("news", []byte("Hello"))
if count != 2 {
t.Errorf("Publish() = %d, want 2", count)
}
}
func TestHubPublishEmptyChannel(t *testing.T) {
hub := NewHub()
count := hub.Publish("empty", []byte("Hello"))
if count != 0 {
t.Errorf("Publish() to empty channel = %d, want 0", count)
}
}
func TestHubChannelCount(t *testing.T) {
hub := NewHub()
client := NewMockClient()
if hub.ChannelCount("news") != 0 {
t.Error("new channel should have 0 subscribers")
}
hub.Subscribe(client, "news")
if hub.ChannelCount("news") != 1 {
t.Errorf("ChannelCount() = %d, want 1", hub.ChannelCount("news"))
}
hub.Subscribe(client, "news") // duplicate
if hub.ChannelCount("news") != 1 {
t.Errorf("ChannelCount() after duplicate = %d, want 1", hub.ChannelCount("news"))
}
}
func TestHubTotalChannels(t *testing.T) {
hub := NewHub()
client := NewMockClient()
if hub.TotalChannels() != 0 {
t.Errorf("TotalChannels() = %d, want 0", hub.TotalChannels())
}
hub.Subscribe(client, "news")
hub.Subscribe(client, "sports")
if hub.TotalChannels() != 2 {
t.Errorf("TotalChannels() = %d, want 2", hub.TotalChannels())
}
}
func TestHubChannelList(t *testing.T) {
hub := NewHub()
client := NewMockClient()
channels := hub.ChannelList("")
if len(channels) != 0 {
t.Errorf("Empty hub ChannelList() length = %d, want 0", len(channels))
}
hub.Subscribe(client, "news")
hub.Subscribe(client, "sports")
channels = hub.ChannelList("")
if len(channels) != 2 {
t.Errorf("ChannelList() length = %d, want 2", len(channels))
}
}
func TestHubMultipleClients(t *testing.T) {
hub := NewHub()
client1 := NewMockClient()
client2 := NewMockClient()
client3 := NewMockClient()
hub.Subscribe(client1, "news")
hub.Subscribe(client2, "news")
hub.Subscribe(client3, "news")
count := hub.ChannelCount("news")
if count != 3 {
t.Errorf("ChannelCount() = %d, want 3", count)
}
}
func TestHubChannelCleanup(t *testing.T) {
hub := NewHub()
client := NewMockClient()
hub.Subscribe(client, "temp")
if hub.ChannelCount("temp") != 1 {
t.Error("Channel should exist")
}
hub.Unsubscribe(client, "temp")
if hub.ChannelCount("temp") != 0 {
t.Error("Channel should be cleaned up")
}
channels := hub.ChannelList("")
for _, ch := range channels {
if ch == "temp" {
t.Error("Channel should not be in list after cleanup")
}
}
}
func TestNewClient(t *testing.T) {
conn := &mockConn{data: []byte{}}
client := NewClient(conn)
if client == nil {
t.Fatal("NewClient() returned nil")
}
if client.subscribed == nil {
t.Error("client.subscribed is nil")
}
}
func TestClientAddSubscription(t *testing.T) {
client := NewMockClient()
client.AddSubscription("news")
if client.SubsCount() != 1 {
t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
}
client.AddSubscription("sports")
if client.SubsCount() != 2 {
t.Errorf("SubsCount() = %d, want 2", client.SubsCount())
}
}
func TestClientRemoveSubscription(t *testing.T) {
client := NewMockClient()
client.AddSubscription("news")
client.AddSubscription("sports")
client.RemoveSubscription("news")
if client.SubsCount() != 1 {
t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
}
}
func TestClientGetChannels(t *testing.T) {
client := NewMockClient()
client.AddSubscription("news")
client.AddSubscription("sports")
channels := client.GetChannels()
if len(channels) != 2 {
t.Errorf("GetChannels() length = %d, want 2", len(channels))
}
found := make(map[string]bool)
for _, ch := range channels {
found[ch] = true
}
if !found["news"] || !found["sports"] {
t.Error("GetChannels() missing expected channels")
}
}
func TestClientGetChannelsReturnsCopy(t *testing.T) {
client := NewMockClient()
client.AddSubscription("news")
channels := client.GetChannels()
channels[0] = "modified"
channels2 := client.GetChannels()
if channels2[0] == "modified" {
t.Error("GetChannels() should return a copy, not modify original")
}
}
func TestClientSubsCount(t *testing.T) {
client := NewMockClient()
if client.SubsCount() != 0 {
t.Errorf("Initial SubsCount() = %d, want 0", client.SubsCount())
}
client.AddSubscription("news")
if client.SubsCount() != 1 {
t.Errorf("SubsCount() = %d, want 1", client.SubsCount())
}
}
func TestConcurrentSubscribe(t *testing.T) {
hub := NewHub()
client := NewMockClient()
done := make(chan bool)
go func() {
hub.Subscribe(client, "channel1")
done <- true
}()
hub.Subscribe(client, "channel2")
<-done
if hub.TotalChannels() != 2 {
t.Errorf("TotalChannels() = %d, want 2", hub.TotalChannels())
}
}

125
internal/server/conn.go Normal file
View File

@@ -0,0 +1,125 @@
package server
import (
"bufio"
"fmt"
"net"
"github.com/eleksir/pekoe/internal/protocol"
)
// Command represents a parsed Redis command with arguments.
type Command struct {
Args [][]byte
}
// ConnReader handles reading and parsing RESP commands from a connection.
// Uses bufio.Reader for efficient I/O with appropriate buffer size.
type ConnReader struct {
reader *bufio.Reader
maxBytes int
}
// NewConnReader creates a new connection reader with the specified max message size.
func NewConnReader(conn net.Conn, maxBytes int) *ConnReader {
return &ConnReader{
reader: bufio.NewReaderSize(conn, 16*1024),
maxBytes: maxBytes,
}
}
// ReadCommand reads and parses a complete Redis command from the connection.
// Returns a Command ready for processing, or an error if parsing fails.
// nolint: cyclop
func (r *ConnReader) ReadCommand() (*Command, error) {
val, err := protocol.Parse(r.reader)
if err != nil {
return nil, err
}
if val == nil {
return nil, fmt.Errorf("unexpected null value")
}
// RESP arrays represent commands - the first element is the command name
array, ok := val.(protocol.Array)
if !ok {
return nil, fmt.Errorf("expected array, got %T", val)
}
if len(array) == 0 {
return nil, fmt.Errorf("empty command array")
}
// Convert array elements to byte slices for command processing
cmd := &Command{
Args: make([][]byte, len(array)),
}
for i, v := range array {
var bytes []byte
switch item := v.(type) {
case protocol.Bulk:
bytes = item
case []byte:
bytes = item
case string:
bytes = []byte(item)
default:
return nil, fmt.Errorf("invalid argument type: %T", v)
}
// Enforce max message size to prevent memory exhaustion
if len(bytes) > r.maxBytes {
return nil, fmt.Errorf("argument too large: %d bytes (max: %d)", len(bytes), r.maxBytes)
}
cmd.Args[i] = bytes
}
return cmd, nil
}
// ConnWriter handles writing RESP responses to a connection.
// Uses buffered writer for efficiency - reduces syscalls.
type ConnWriter struct {
conn net.Conn
writer *bufio.Writer
}
// NewConnWriter creates a new connection writer with buffering.
func NewConnWriter(conn net.Conn) *ConnWriter {
return &ConnWriter{
conn: conn,
writer: bufio.NewWriterSize(conn, 16*1024),
}
}
// WriteStatus writes a simple status response (+OK).
func (w *ConnWriter) WriteStatus(status string) {
_, _ = w.writer.Write(protocol.MakeStatus(status))
}
// WriteError writes an error response (-ERR message).
func (w *ConnWriter) WriteError(err string) {
_, _ = w.writer.Write(protocol.MakeError(err))
}
// WriteInt writes an integer response (:12345).
func (w *ConnWriter) WriteInt(n int64) {
_, _ = w.writer.Write(protocol.MakeInt(n))
}
// WriteBulk writes a bulk string response ($<len>\r\n<data>).
func (w *ConnWriter) WriteBulk(data []byte) {
_, _ = w.writer.Write(protocol.MakeBulk(data))
}
// WriteArray writes a RESP array response - used for pub/sub confirmations.
func (w *ConnWriter) WriteArray(values []protocol.Value) {
_, _ = w.writer.Write(protocol.MakeMultiBulk(values))
}
// Flush writes all buffered data to the underlying connection.
func (w *ConnWriter) Flush() error {
return w.writer.Flush()
}

479
internal/server/handlers.go Normal file
View File

@@ -0,0 +1,479 @@
// @tag server
// @tag handlers:pubsub
package server
import (
"fmt"
"strings"
"github.com/eleksir/pekoe/internal/logger"
"github.com/eleksir/pekoe/internal/protocol"
"github.com/eleksir/pekoe/internal/pubsub"
)
// handleSubscribe processes SUBSCRIBE command.
// SUBSCRIBE channel [channel ...] - Subscribe to one or more channels.
// Returns immediately after subscription (doesn't block like PSUBSCRIBE).
func (s *Server) handleSubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error {
if len(channels) == 0 {
writer.WriteError("ERR wrong number of arguments for 'subscribe' command")
return nil
}
// Subscribe to each channel specified
for _, ch := range channels {
channel := string(ch)
if s.hub.Subscribe(client, channel) {
logger.Debug("client=%s event=subscribe channel=%s", client.RemoteAddr(), channel)
// Send confirmation in format: ["subscribe", channel, count]
writer.WriteArray([]protocol.Value{
protocol.Bulk("subscribe"),
protocol.Bulk(channel),
protocol.Int(int64(client.SubsCount())),
})
}
}
return nil
}
// handleUnsubscribe processes UNSUBSCRIBE command.
// UNSUBSCRIBE [channel [channel ...]] - Unsubscribe from channels.
// If no channels specified, unsubscribe from all.
func (s *Server) handleUnsubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error {
if len(channels) == 0 {
// Unsubscribe from all channels
chans := client.GetChannels()
if len(chans) == 0 {
// No subscriptions - send single unsubscribe message with -1 (Redis behavior)
writer.WriteArray([]protocol.Value{
protocol.Bulk("unsubscribe"),
protocol.Bulk("-1"),
protocol.Int(0),
})
return nil
}
for _, ch := range chans {
s.hub.Unsubscribe(client, ch)
writer.WriteArray([]protocol.Value{
protocol.Bulk("unsubscribe"),
protocol.Bulk(ch),
protocol.Int(int64(client.SubsCount())),
})
}
} else {
// Unsubscribe from specified channels - only send response if actually subscribed
for _, ch := range channels {
channel := string(ch)
removed := s.hub.Unsubscribe(client, channel)
if removed {
logger.Debug("client=%s event=unsubscribe channel=%s", client.RemoteAddr(), channel)
writer.WriteArray([]protocol.Value{
protocol.Bulk("unsubscribe"),
protocol.Bulk(channel),
protocol.Int(int64(client.SubsCount())),
})
}
}
}
return nil
}
// handlePublish processes PUBLISH command.
// PUBLISH channel message - Publish a message to a channel.
// Returns the number of subscribers who received the message.
func (s *Server) handlePublish(args [][]byte, writer *ConnWriter) error {
if len(args) != 2 {
writer.WriteError("ERR wrong number of arguments for 'publish' command")
return nil
}
channel := string(args[0])
message := args[1]
regularCount := s.hub.Publish(channel, message)
patternCount := s.hub.PublishToPattern(channel, message)
totalCount := regularCount + patternCount
logger.Debug("event=publish channel=%s regular=%d pattern=%d", channel, regularCount, patternCount)
writer.WriteInt(int64(totalCount))
return nil
}
// handlePubsub processes PUBSUB subcommands for introspection.
// nolint: cyclop
func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error {
if len(args) == 0 {
writer.WriteError("ERR wrong number of arguments for 'pubsub' command")
return nil
}
subcmd := strings.ToLower(string(args[0]))
switch subcmd {
case "numsub":
// PUBSUB NUMSUB [channel ...] - Get subscriber count for channels
if len(args) < 2 {
writer.WriteError("ERR wrong number of arguments for 'pubsub numsub' command")
return nil
}
// Build response: [[channel, count], ...]
channels := args[1:]
result := make([]protocol.Value, 0, len(channels)*2)
for _, ch := range channels {
channel := string(ch)
count := s.hub.ChannelCount(channel)
result = append(result, protocol.Bulk(ch), protocol.Int(int64(count)))
}
writer.WriteArray(result)
case "channels":
// PUBSUB CHANNELS [pattern] - List active channels, optionally filtered
var pattern string
if len(args) > 1 {
pattern = string(args[1])
}
channels := s.channelList(pattern)
result := make([]protocol.Value, len(channels))
for i, ch := range channels {
result[i] = protocol.Bulk(ch)
}
writer.WriteArray(result)
case "numpat":
// PUBSUB NUMPAT - Count of unique pattern subscriptions
writer.WriteInt(int64(s.hub.PatternCount()))
case "help":
// PUBSUB HELP - Returns helpful text about pubsub subcommands
helpText := []string{
"PUBSUB <subcommand> [<argument> [<argument> ...]]>",
" PUBSUB CHANNELS [pattern] - List active channels (optionally filtered by pattern)",
" PUBSUB NUMSUB [channel ...] - Get subscriber count for channels",
" PUBSUB NUMPAT - Get the number of pattern subscriptions",
" PUBSUB HELP - Show this help message",
}
result := make([]protocol.Value, len(helpText))
for i, line := range helpText {
result[i] = protocol.Bulk(line)
}
writer.WriteArray(result)
default:
writer.WriteError(fmt.Sprintf("ERR unknown pubsub subcommand '%s'", subcmd))
}
return nil
}
// channelList returns all active channel names, optionally filtered by pattern.
func (s *Server) channelList(pattern string) []string {
channels := s.hub.ChannelList(pattern)
if channels == nil {
return []string{}
}
return channels
}
// handlePing processes PING command.
// PING [message] - Returns PONG or the optional message.
// Redis returns simple string for message, not bulk string.
func (s *Server) handlePing(args [][]byte, writer *ConnWriter) error {
if len(args) > 0 {
writer.WriteStatus(string(args[0]))
} else {
writer.WriteStatus("PONG")
}
return nil
}
// handleQuit processes QUIT command.
// QUIT - Close the connection gracefully.
func (s *Server) handleQuit(client *pubsub.Client, writer *ConnWriter) error {
writer.WriteStatus("OK")
return ErrQuit
}
// handleClient processes CLIENT command.
// CLIENT subcommand [argument [argument ...]]
// nolint: cyclop
func (s *Server) handleClient(client *pubsub.Client, args [][]byte, writer *ConnWriter) error {
if len(args) == 0 {
writer.WriteError("ERR wrong number of arguments for 'client' command")
return nil
}
subcmd := strings.ToLower(string(args[0]))
switch subcmd {
case "id":
// CLIENT ID - return the connection ID
//nolint: gosec // client ID is sequential, overflow is theoretical
writer.WriteInt(int64(client.ID()))
case "setname":
// CLIENT SETNAME connection-name
if len(args) < 2 {
writer.WriteError("ERR wrong number of arguments for 'client setname' command")
return nil
}
name := string(args[1])
// Validate name - must be printable characters and no spaces
if !isValidClientName(name) {
writer.WriteError("ERR Client names cannot contain spaces, null characters or names longer than 255 characters")
return nil
}
client.SetName(name)
writer.WriteStatus("OK")
case "getname":
// CLIENT GETNAME - return the connection name
name := client.GetName()
if name == "" {
writer.WriteBulk(nil)
} else {
writer.WriteBulk([]byte(name))
}
case "list":
// CLIENT LIST [TYPE normal|master|slave|pubsub] [ID min-id max-id]
// For pekoe, we only support pubsub clients
clientList := s.buildClientList(args[1:])
writer.WriteBulk([]byte(clientList))
case "info":
// CLIENT INFO - return information about the client
writer.WriteBulk([]byte(client.InfoString()))
case "kill":
// CLIENT KILL [ID min-id max-id] [ADDR ip:port] [LADDR ip:port] [TYPE normal|master|slave|pubsub] [USER username] [SKIPME yes|no]
count, err := s.handleClientKill(client, args[1:])
if err != nil {
writer.WriteError(err.Error())
return nil
}
writer.WriteInt(int64(count))
case "pause":
// CLIENT PAUSE timeout [WRITE|ALL]
// Acknowledge but don't actually implement - pubsub doesn't need this
writer.WriteStatus("OK")
case "unpause":
// CLIENT UNPAUSE
writer.WriteStatus("OK")
case "reply":
// CLIENT REPLY ON|OFF|SKIP
if len(args) < 2 {
writer.WriteError("ERR wrong number of arguments for 'client reply' command")
return nil
}
mode := strings.ToUpper(string(args[1]))
switch mode {
case "ON":
client.SetReplyMode(pubsub.ReplyOn)
writer.WriteStatus("OK")
case "OFF":
client.SetReplyMode(pubsub.ReplyOff)
// Don't send reply when setting to OFF
_ = writer.Flush()
case "SKIP":
client.SetReplyMode(pubsub.ReplySkip)
// Send one skip reply, then next will be skipped
writer.WriteStatus("OK")
default:
writer.WriteError("ERR wrong argument for 'client reply' command")
}
case "tracking":
// CLIENT TRACKING ON/OFF [REDIRECT client-id] [PREFIX prefix ...] [BCAST] [OPTIN] [OPTOUT] [NOLOOP]
// Acknowledge but don't implement - pekoe is pub/sub only
writer.WriteStatus("OK")
case "no-evict":
// CLIENT NO-EVICT ON|OFF
writer.WriteStatus("OK")
case "no-touch":
// CLIENT NO-TOUCH ON|OFF
writer.WriteStatus("OK")
case "caching":
// CLIENT CACHING YES|NO
writer.WriteStatus("OK")
default:
// Unknown client subcommand
writer.WriteError(fmt.Sprintf("ERR Unknown subcommand or wrong number of arguments for '%s'. Try CLIENT HELP.", subcmd))
}
return nil
}
// isValidClientName checks if a client name is valid per Redis spec
func isValidClientName(name string) bool {
if len(name) == 0 || len(name) > 255 {
return false
}
for _, c := range name {
if c < 32 || c > 126 || c == ' ' {
return false
}
}
return true
}
// buildClientList builds CLIENT LIST response for all or filtered clients
func (s *Server) buildClientList(args [][]byte) string {
clients := s.hub.GetAllClients()
// Parse optional filter arguments
var filterType string
for i := 0; i < len(args); i++ {
arg := strings.ToLower(string(args[i]))
if arg == "type" && i+1 < len(args) {
filterType = strings.ToLower(string(args[i+1]))
i++
}
}
var result string
for _, c := range clients {
// For pekoe, all clients are "pubsub" type
if filterType != "" && filterType != "pubsub" {
continue
}
result += c.ListString() + "\n"
}
return result
}
// handleClientKill handles CLIENT KILL subcommand
// nolint: cyclop
func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (int, error) {
clients := s.hub.GetAllClients()
killed := 0
// Parse arguments
var killID uint64
var killAddr string
var skipMe = true
for i := 0; i < len(args); i++ {
arg := strings.ToUpper(string(args[i]))
switch arg {
case "ID":
if i+2 < len(args) {
// ID min max format - kill all IDs in range
i += 2
} else if i+1 < len(args) {
// Single ID
if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killID); err != nil {
return 0, fmt.Errorf("invalid client ID: %s", args[i+1])
}
i++
}
case "ADDR":
if i+1 < len(args) {
killAddr = string(args[i+1])
i++
}
case "SKIPME":
if i+1 < len(args) {
skipMe = strings.ToUpper(string(args[i+1])) == "YES"
i++
}
}
}
for _, c := range clients {
// Skip current client if requested
if skipMe && c.ID() == currentClient.ID() {
continue
}
shouldKill := false
if killID > 0 && c.ID() == killID {
shouldKill = true
}
if killAddr != "" && c.RemoteAddr() == killAddr {
shouldKill = true
}
if shouldKill {
// Unsubscribe and close
s.hub.UnsubscribeAll(c)
s.hub.UnregisterClient(c)
_ = c.Close()
killed++
}
}
return killed, 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 {
if len(patterns) == 0 {
writer.WriteError("ERR wrong number of arguments for 'psubscribe' command")
return nil
}
for _, p := range patterns {
pattern := string(p)
ok, err := s.hub.PatternSubscribe(client, pattern)
if err != nil {
writer.WriteError(fmt.Sprintf("ERR invalid pattern: %v", err))
return nil
}
if ok {
logger.Debug("client=%s event=psubscribe pattern=%s", client.RemoteAddr(), pattern)
// Send confirmation in format: ["psubscribe", pattern, count]
writer.WriteArray([]protocol.Value{
protocol.Bulk("psubscribe"),
protocol.Bulk(pattern),
protocol.Int(int64(client.PatternSubsCount())),
})
}
}
return nil
}
// handlePunsubscribe processes PUNSUBSCRIBE command.
// PUNSUBSCRIBE [pattern [pattern ...]] - Unsubscribe from patterns.
func (s *Server) handlePunsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error {
if len(patterns) == 0 {
// Unsubscribe from all patterns
pats := client.GetPatterns()
if len(pats) == 0 {
// No pattern subscriptions - send single punsubscribe message with -1 (Redis behavior)
writer.WriteArray([]protocol.Value{
protocol.Bulk("punsubscribe"),
protocol.Bulk("-1"),
protocol.Int(0),
})
return nil
}
for _, pattern := range pats {
s.hub.PatternUnsubscribe(client, pattern)
writer.WriteArray([]protocol.Value{
protocol.Bulk("punsubscribe"),
protocol.Bulk(pattern),
protocol.Int(int64(client.PatternSubsCount())),
})
}
} else {
// Unsubscribe from specified patterns - only send response if actually subscribed
for _, p := range patterns {
pattern := string(p)
removed := s.hub.PatternUnsubscribe(client, pattern)
if removed {
logger.Debug("client=%s event=punsubscribe pattern=%s", client.RemoteAddr(), pattern)
writer.WriteArray([]protocol.Value{
protocol.Bulk("punsubscribe"),
protocol.Bulk(pattern),
protocol.Int(int64(client.PatternSubsCount())),
})
}
}
}
return nil
}

325
internal/server/server.go Normal file
View 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
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,46 @@
//go:build !windows
package server
import (
"os"
"os/signal"
"syscall"
"github.com/eleksir/pekoe/internal/logger"
)
// signalSetup configures signal handlers for Unix systems.
// Returns a channel that receives incoming signals.
func signalSetup() chan os.Signal {
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGUSR1)
return sigCh
}
// handleSignal processes a received signal and returns true if the server should shutdown.
func (s *Server) handleSignal(sig os.Signal) bool {
switch sig {
case syscall.SIGUSR1:
// Reopen log files for log rotation
if err := logger.Reopen(); err != nil {
logger.Error("event=log_reopen_error error=%v", err)
} else {
logger.Info("event=log_reopened")
}
return false // Continue listening
case syscall.SIGINT, syscall.SIGTERM:
// Graceful shutdown
logger.Info("event=shutdown signal=%v", sig)
return true // Shutdown
case syscall.SIGHUP:
// Reserved for future config reload functionality
logger.Info("event=sighup_ignored message=reload not implemented")
return false // Continue listening
default:
return false
}
}

View File

@@ -0,0 +1,43 @@
//go:build windows
package server
import (
"os"
"os/signal"
"syscall"
"github.com/eleksir/pekoe/internal/logger"
)
// signalSetup configures signal handlers for Windows.
// Windows doesn't support SIGUSR1, so we only handle shutdown signals.
// For log rotation on Windows, use a different mechanism.
func signalSetup() chan os.Signal {
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
return sigCh
}
// handleSignal processes a received signal and returns true if the server should shutdown.
func (s *Server) handleSignal(sig os.Signal) bool {
switch sig {
case syscall.SIGINT, syscall.SIGTERM:
// Graceful shutdown
logger.Info("event=shutdown signal=%v", sig)
return true // Shutdown
case syscall.SIGHUP:
// On Windows, SIGHUP can be used for log rotation since SIGUSR1 is not available
// Reopen log files
if err := logger.Reopen(); err != nil {
logger.Error("event=log_reopen_error error=%v", err)
} else {
logger.Info("event=log_reopened")
}
return false // Continue listening
default:
return false
}
}