From b68ae5e548116c0f7da002774c9beba0365513f7 Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Thu, 13 Nov 2014 14:26:14 +0200
Subject: [PATCH] Fix rate limiter and add test.

---
 Makefile           |  2 +-
 commands.go        | 12 +++++++++++-
 rate_limit.go      |  3 ++-
 rate_limit_test.go | 28 ++++++++++++++++++++++++++++
 sentinel_test.go   | 26 +++++++++++++++++++-------
 5 files changed, 61 insertions(+), 10 deletions(-)
 create mode 100644 rate_limit_test.go

diff --git a/Makefile b/Makefile
index f9571a4..6d6d427 100644
--- a/Makefile
+++ b/Makefile
@@ -1,2 +1,2 @@
 all:
-	go test gopkg.in/redis.v2
+	go test gopkg.in/redis.v2 -cpu=1,2,4
diff --git a/commands.go b/commands.go
index 56cdf9a..1fc22aa 100644
--- a/commands.go
+++ b/commands.go
@@ -1,6 +1,7 @@
 package redis
 
 import (
+	"io"
 	"strconv"
 	"time"
 )
@@ -1120,7 +1121,16 @@ func (c *Client) shutdown(modifier string) *StatusCmd {
 	}
 	cmd := NewStatusCmd(args...)
 	c.Process(cmd)
-	c.Close()
+	if err := cmd.Err(); err != nil {
+		if err == io.EOF {
+			// Server quit as expected.
+			cmd.err = nil
+		}
+	} else {
+		// Server did not quit. String reply contains the reason.
+		cmd.err = errorf(cmd.val)
+		cmd.val = ""
+	}
 	return cmd
 }
 
diff --git a/rate_limit.go b/rate_limit.go
index a046648..20d8512 100644
--- a/rate_limit.go
+++ b/rate_limit.go
@@ -37,8 +37,9 @@ func (rl *rateLimiter) Check() bool {
 			if atomic.CompareAndSwapInt64(&rl.v, v, v-1) {
 				return true
 			}
+		} else {
+			return false
 		}
-		return false
 	}
 }
 
diff --git a/rate_limit_test.go b/rate_limit_test.go
new file mode 100644
index 0000000..f6a4721
--- /dev/null
+++ b/rate_limit_test.go
@@ -0,0 +1,28 @@
+package redis
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestRateLimiter(t *testing.T) {
+	const n = 100000
+	rl := newRateLimiter(time.Second, n)
+
+	wg := &sync.WaitGroup{}
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func() {
+			if !rl.Check() {
+				panic("check failed")
+			}
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	if rl.Check() && rl.Check() {
+		t.Fatal("check passed")
+	}
+}
diff --git a/sentinel_test.go b/sentinel_test.go
index ae317b7..ede59bd 100644
--- a/sentinel_test.go
+++ b/sentinel_test.go
@@ -88,10 +88,13 @@ func TestSentinel(t *testing.T) {
 	slavePort := "8124"
 	sentinelPort := "8125"
 
-	_, err := startRedis(masterPort)
+	masterCmd, err := startRedis(masterPort)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer masterCmd.Process.Kill()
+
+	// Wait for master to start.
 	time.Sleep(200 * time.Millisecond)
 
 	master := redis.NewTCPClient(&redis.Options{
@@ -100,12 +103,14 @@ func TestSentinel(t *testing.T) {
 	if err := master.Ping().Err(); err != nil {
 		t.Fatal(err)
 	}
-	defer master.Shutdown()
 
-	_, err = startRedisSlave(slavePort, masterPort)
+	slaveCmd, err := startRedisSlave(slavePort, masterPort)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer slaveCmd.Process.Kill()
+
+	// Wait for slave to start.
 	time.Sleep(200 * time.Millisecond)
 
 	slave := redis.NewTCPClient(&redis.Options{
@@ -114,12 +119,14 @@ func TestSentinel(t *testing.T) {
 	if err := slave.Ping().Err(); err != nil {
 		t.Fatal(err)
 	}
-	defer slave.Shutdown()
 
-	_, err = startRedisSentinel(sentinelPort, masterName, masterPort)
+	sentinelCmd, err := startRedisSentinel(sentinelPort, masterName, masterPort)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer sentinelCmd.Process.Kill()
+
+	// Wait for sentinel to start.
 	time.Sleep(200 * time.Millisecond)
 
 	sentinel := redis.NewTCPClient(&redis.Options{
@@ -147,8 +154,13 @@ func TestSentinel(t *testing.T) {
 		t.Fatalf(`got %q, expected "master"`, val)
 	}
 
-	// Kill master.
-	master.Shutdown()
+	// Kill Redis master.
+	if err := masterCmd.Process.Kill(); err != nil {
+		t.Fatal(err)
+	}
+	if err := master.Ping().Err(); err == nil {
+		t.Fatalf("master was not killed")
+	}
 
 	// Wait for Redis sentinel to elect new master.
 	time.Sleep(5 * time.Second)