Initial commit
This commit is contained in:
959
internal/compatibility/goredis_test.go
Normal file
959
internal/compatibility/goredis_test.go
Normal 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
|
||||
}
|
||||
118
internal/compatibility/helper_test.go
Normal file
118
internal/compatibility/helper_test.go
Normal 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
|
||||
}
|
||||
}
|
||||
833
internal/compatibility/redigo_test.go
Normal file
833
internal/compatibility/redigo_test.go
Normal 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)
|
||||
}
|
||||
217
internal/compatibility/rueidis_test.go
Normal file
217
internal/compatibility/rueidis_test.go
Normal 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
169
internal/config/config.go
Normal 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)
|
||||
}
|
||||
435
internal/config/config_test.go
Normal file
435
internal/config/config_test.go
Normal 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")
|
||||
}
|
||||
}
|
||||
24
internal/config/configpath_bsd.go
Normal file
24
internal/config/configpath_bsd.go
Normal 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",
|
||||
}
|
||||
}
|
||||
29
internal/config/configpath_linux.go
Normal file
29
internal/config/configpath_linux.go
Normal 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",
|
||||
}
|
||||
}
|
||||
11
internal/config/configpath_windows.go
Normal file
11
internal/config/configpath_windows.go
Normal 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
170
internal/logger/logger.go
Normal 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)
|
||||
}
|
||||
359
internal/logger/logger_test.go
Normal file
359
internal/logger/logger_test.go
Normal 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()
|
||||
}
|
||||
252
internal/protocol/protocol.go
Normal file
252
internal/protocol/protocol.go
Normal 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)))
|
||||
}
|
||||
}
|
||||
206
internal/protocol/protocol_test.go
Normal file
206
internal/protocol/protocol_test.go
Normal 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
257
internal/pubsub/client.go
Normal 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
463
internal/pubsub/hub.go
Normal 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
|
||||
}
|
||||
333
internal/pubsub/pubsub_test.go
Normal file
333
internal/pubsub/pubsub_test.go
Normal 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
125
internal/server/conn.go
Normal 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
479
internal/server/handlers.go
Normal 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
325
internal/server/server.go
Normal file
@@ -0,0 +1,325 @@
|
||||
// @tag server
|
||||
// @tag core
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/eleksir/pekoe/internal/config"
|
||||
"github.com/eleksir/pekoe/internal/logger"
|
||||
"github.com/eleksir/pekoe/internal/pubsub"
|
||||
)
|
||||
|
||||
// Sentinel errors for connection handling
|
||||
var (
|
||||
ErrQuit = errors.New("client quit")
|
||||
)
|
||||
|
||||
// Server is the main Pekoe server instance.
|
||||
// It manages the TCP listener, connection lifecycle, and delegates to the pubsub hub.
|
||||
type Server struct {
|
||||
cfg *config.Config
|
||||
hub pubsubHub
|
||||
ln net.Listener
|
||||
wg sync.WaitGroup
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
active atomic.Int64 // atomic counter for active connections
|
||||
}
|
||||
|
||||
// PubSubHub defines the interface our server expects from the pubsub implementation.
|
||||
// This allows for cleaner dependency injection and testing.
|
||||
type pubsubHub interface {
|
||||
Subscribe(client *pubsub.Client, channel string) bool
|
||||
Unsubscribe(client *pubsub.Client, channel string) bool
|
||||
UnsubscribeAll(client *pubsub.Client)
|
||||
Publish(channel string, message []byte) int
|
||||
PublishToPattern(channel string, message []byte) int
|
||||
ChannelCount(channel string) int
|
||||
ChannelList(pattern string) []string
|
||||
TotalChannels() int
|
||||
PatternSubscribe(client *pubsub.Client, pattern string) (bool, error)
|
||||
PatternUnsubscribe(client *pubsub.Client, pattern string) bool
|
||||
PatternCount() int
|
||||
RegisterClient(client *pubsub.Client)
|
||||
UnregisterClient(client *pubsub.Client)
|
||||
GetClient(id uint64) *pubsub.Client
|
||||
GetAllClients() []*pubsub.Client
|
||||
ClientCount() int
|
||||
}
|
||||
|
||||
// New creates a new Pekoe server with the given config and pubsub hub.
|
||||
func New(cfg *config.Config, hub pubsubHub) *Server {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
hub: hub,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
// Run starts the server and blocks until shutdown.
|
||||
// Listens on the configured address, accepts connections, and handles them.
|
||||
// Returns an error only if the initial listen fails.
|
||||
func (s *Server) Run() error {
|
||||
ln, err := net.Listen("tcp", s.cfg.Addr())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to listen on %s: %w", s.cfg.Addr(), err)
|
||||
}
|
||||
s.ln = ln
|
||||
|
||||
logger.Info("event=startup addr=%s max_clients=%d", s.cfg.Addr(), s.cfg.MaxClients)
|
||||
|
||||
// Create PID file if configured
|
||||
if s.cfg.PidFile != "" {
|
||||
if err := s.createPidFile(); err != nil {
|
||||
logger.Error("event=pid_file_error error=%v", err)
|
||||
// Don't fail startup, just log the error
|
||||
} else {
|
||||
logger.Info("event=pid_file_created path=%s pid=%d", s.cfg.PidFile, os.Getpid())
|
||||
}
|
||||
}
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.acceptLoop()
|
||||
|
||||
// Wait for shutdown signal
|
||||
s.waitForSignal()
|
||||
return nil
|
||||
}
|
||||
|
||||
// createPidFile creates the PID file with the current process ID
|
||||
func (s *Server) createPidFile() error {
|
||||
pid := strconv.Itoa(os.Getpid())
|
||||
//nolint: gosec
|
||||
return os.WriteFile(s.cfg.PidFile, []byte(pid), 0644)
|
||||
}
|
||||
|
||||
// removePidFile removes the PID file on shutdown
|
||||
func (s *Server) removePidFile() {
|
||||
if s.cfg.PidFile == "" {
|
||||
return
|
||||
}
|
||||
|
||||
if err := os.Remove(s.cfg.PidFile); err != nil {
|
||||
if !errors.Is(err, fs.ErrNotExist) {
|
||||
logger.Error("event=pid_file_remove_error path=%s error=%v", s.cfg.PidFile, err)
|
||||
}
|
||||
} else {
|
||||
logger.Info("event=pid_file_removed path=%s", s.cfg.PidFile)
|
||||
}
|
||||
}
|
||||
|
||||
// AcceptLoop runs in a goroutine and accepts incoming connections.
|
||||
// Uses a loop with error handling to continuously accept connections until shutdown.
|
||||
func (s *Server) acceptLoop() {
|
||||
defer s.wg.Done()
|
||||
|
||||
for {
|
||||
conn, err := s.ln.Accept()
|
||||
if err != nil {
|
||||
// Check if we're shutting down (context cancelled) vs real error
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Handle temporary errors gracefully - they're often transient
|
||||
//nolint: staticcheck
|
||||
if ne, ok := err.(net.Error); ok && ne.Temporary() {
|
||||
logger.Warn("event=accept_temporary error=%v retry_ms=5", err)
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
continue
|
||||
}
|
||||
logger.Error("event=accept_error error=%v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Enforce max clients limit - protects server from resource exhaustion
|
||||
if s.active.Load() >= int64(s.cfg.MaxClients) {
|
||||
logger.Warn("event=max_clients_reached client=%s", conn.RemoteAddr())
|
||||
_ = conn.Close()
|
||||
continue
|
||||
}
|
||||
|
||||
go s.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
// HandleConn manages the lifecycle of a single client connection.
|
||||
// Parses commands, processes them, and sends responses.
|
||||
// Uses per-connection goroutine to allow concurrent client handling.
|
||||
// nolint: cyclop
|
||||
func (s *Server) handleConn(conn net.Conn) {
|
||||
s.wg.Add(1)
|
||||
defer s.wg.Done()
|
||||
s.active.Add(1)
|
||||
defer s.active.Add(-1)
|
||||
|
||||
// Create and register client
|
||||
client := pubsub.NewClient(conn)
|
||||
s.hub.RegisterClient(client)
|
||||
defer func() {
|
||||
s.hub.UnregisterClient(client)
|
||||
s.hub.UnsubscribeAll(client)
|
||||
_ = client.Close()
|
||||
}()
|
||||
|
||||
logger.Info("client=%s event=connected client_id=%d", client.RemoteAddr(), client.ID())
|
||||
|
||||
// Set timeouts if configured
|
||||
if s.cfg.ReadTimeoutDuration() > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeoutDuration()))
|
||||
}
|
||||
if s.cfg.WriteTimeoutDuration() > 0 {
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(s.cfg.WriteTimeoutDuration()))
|
||||
}
|
||||
|
||||
reader := NewConnReader(conn, s.cfg.MaxMsgBytes)
|
||||
writer := NewConnWriter(conn)
|
||||
|
||||
// Main command loop - process commands until error or disconnect
|
||||
for {
|
||||
// Check for shutdown before each command
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Reset read deadline before each command if using timeouts
|
||||
if s.cfg.ReadTimeoutDuration() > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeoutDuration()))
|
||||
}
|
||||
|
||||
cmd, err := reader.ReadCommand()
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
logger.Info("client=%s event=disconnected client_id=%d", client.RemoteAddr(), client.ID())
|
||||
return // Client disconnected normally
|
||||
}
|
||||
logger.Error("client=%s error=protocol message=%v", client.RemoteAddr(), err)
|
||||
writer.WriteError("Protocol error: " + err.Error())
|
||||
_ = writer.Flush()
|
||||
return
|
||||
}
|
||||
|
||||
// Log incoming command at DEBUG level
|
||||
cmdName := ""
|
||||
if len(cmd.Args) > 0 {
|
||||
cmdName = strings.ToLower(string(cmd.Args[0]))
|
||||
}
|
||||
logger.Debug("client=%s command=%s args=%d", client.RemoteAddr(), cmdName, len(cmd.Args)-1)
|
||||
|
||||
// Process command - handler returns error only on fatal issues
|
||||
if err := s.processCommand(client, cmd, writer); err != nil {
|
||||
if errors.Is(err, ErrQuit) {
|
||||
_ = writer.Flush()
|
||||
logger.Info("client=%s event=quit", client.RemoteAddr())
|
||||
return // Graceful client disconnect via QUIT
|
||||
}
|
||||
logger.Error("client=%s error=processing message=%v", client.RemoteAddr(), err)
|
||||
return
|
||||
}
|
||||
|
||||
// Flush response - important for pubsub as we need immediate feedback
|
||||
if s.cfg.WriteTimeoutDuration() > 0 {
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(s.cfg.WriteTimeoutDuration()))
|
||||
}
|
||||
if err := writer.Flush(); err != nil {
|
||||
logger.Error("client=%s error=flush message=%v", client.RemoteAddr(), err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processCommand routes commands to their handlers.
|
||||
// Each command type has its own handler for clarity.
|
||||
// nolint: cyclop
|
||||
func (s *Server) processCommand(client *pubsub.Client, cmd *Command, writer *ConnWriter) error {
|
||||
if cmd == nil || len(cmd.Args) == 0 {
|
||||
writer.WriteError("ERR empty command")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Convert command name to lowercase for case-insensitive matching (Redis standard)
|
||||
name := strings.ToLower(string(cmd.Args[0]))
|
||||
switch name {
|
||||
case "subscribe":
|
||||
return s.handleSubscribe(client, cmd.Args[1:], writer)
|
||||
case "unsubscribe":
|
||||
return s.handleUnsubscribe(client, cmd.Args[1:], writer)
|
||||
case "psubscribe":
|
||||
return s.handlePsubscribe(client, cmd.Args[1:], writer)
|
||||
case "punsubscribe":
|
||||
return s.handlePunsubscribe(client, cmd.Args[1:], writer)
|
||||
case "publish":
|
||||
return s.handlePublish(cmd.Args[1:], writer)
|
||||
case "pubsub":
|
||||
return s.handlePubsub(cmd.Args[1:], writer)
|
||||
case "ping":
|
||||
return s.handlePing(cmd.Args[1:], writer)
|
||||
case "quit":
|
||||
return s.handleQuit(client, writer)
|
||||
case "client":
|
||||
return s.handleClient(client, cmd.Args[1:], writer)
|
||||
default:
|
||||
// Check if client is in pubsub mode - only pubsub commands allowed
|
||||
if client != nil && (client.SubsCount() > 0 || client.PatternSubsCount() > 0) {
|
||||
logger.Error("client=%s error=invalid_in_pubsub_mode command=%s", client.RemoteAddr(), name)
|
||||
writer.WriteError("ERR Can't execute commands in pubsub mode")
|
||||
return nil
|
||||
}
|
||||
clientAddr := ""
|
||||
if client != nil {
|
||||
clientAddr = client.RemoteAddr()
|
||||
}
|
||||
logger.Error("client=%s error=unknown_command command=%s", clientAddr, name)
|
||||
writer.WriteError(fmt.Sprintf("ERR unknown command '%s'", name))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// waitForSignal listens for signals and handles them appropriately.
|
||||
// On Unix: handles SIGINT, SIGTERM (shutdown), SIGHUP (reserved), SIGUSR1 (log reopen)
|
||||
// On Windows: handles SIGINT, SIGTERM (shutdown), SIGHUP (log reopen)
|
||||
func (s *Server) waitForSignal() {
|
||||
sigCh := signalSetup()
|
||||
|
||||
for {
|
||||
sig := <-sigCh
|
||||
|
||||
if s.handleSignal(sig) {
|
||||
// Shutdown signal received - perform graceful shutdown
|
||||
s.cancel()
|
||||
|
||||
// Close listener - causes Accept() to return error, stopping accept loop
|
||||
if err := s.ln.Close(); err != nil {
|
||||
logger.Error("event=listener_close_error error=%v", err)
|
||||
}
|
||||
|
||||
// Wait for all active connections to finish
|
||||
s.wg.Wait()
|
||||
|
||||
// Remove PID file
|
||||
s.removePidFile()
|
||||
|
||||
logger.Info("event=shutdown_complete active_connections=%d", s.active.Load())
|
||||
return
|
||||
}
|
||||
// Not a shutdown signal - continue listening
|
||||
}
|
||||
}
|
||||
1046
internal/server/server_test.go
Normal file
1046
internal/server/server_test.go
Normal file
File diff suppressed because it is too large
Load Diff
46
internal/server/signals_unix.go
Normal file
46
internal/server/signals_unix.go
Normal 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
|
||||
}
|
||||
}
|
||||
43
internal/server/signals_windows.go
Normal file
43
internal/server/signals_windows.go
Normal 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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user