From d0d3920e69467352e1ea1b3edf28e2aa2c6e9b4a Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Wed, 4 Nov 2015 14:25:48 +0200
Subject: [PATCH] Make Pipeline thread-safe. Fixes #166.

---
 multi.go         | 27 ++++++++++++++++++++++++---
 pipeline.go      | 36 +++++++++++++++++++++++++++++-------
 pipeline_test.go | 21 +++++++++++++++++++++
 3 files changed, 74 insertions(+), 10 deletions(-)

diff --git a/multi.go b/multi.go
index 00dc9991..d0859e61 100644
--- a/multi.go
+++ b/multi.go
@@ -9,13 +9,18 @@ import (
 var errDiscard = errors.New("redis: Discard can be used only inside Exec")
 
 // Multi implements Redis transactions as described in
-// http://redis.io/topics/transactions. It's NOT safe for concurrent
-// use by multiple goroutines.
+// http://redis.io/topics/transactions. It's NOT safe for concurrent use
+// by multiple goroutines, because Exec resets connection state.
+// If you don't need WATCH it is better to use Pipeline.
+//
+// TODO(vmihailenco): rename to Tx
 type Multi struct {
 	commandable
 
 	base *baseClient
-	cmds []Cmder
+
+	cmds   []Cmder
+	closed bool
 }
 
 func (c *Client) Multi() *Multi {
@@ -37,13 +42,17 @@ func (c *Multi) process(cmd Cmder) {
 	}
 }
 
+// Close closes the client, releasing any open resources.
 func (c *Multi) Close() error {
+	c.closed = true
 	if err := c.Unwatch().Err(); err != nil {
 		log.Printf("redis: Unwatch failed: %s", err)
 	}
 	return c.base.Close()
 }
 
+// Watch marks the keys to be watched for conditional execution
+// of a transaction.
 func (c *Multi) Watch(keys ...string) *StatusCmd {
 	args := make([]interface{}, 1+len(keys))
 	args[0] = "WATCH"
@@ -55,6 +64,7 @@ func (c *Multi) Watch(keys ...string) *StatusCmd {
 	return cmd
 }
 
+// Unwatch flushes all the previously watched keys for a transaction.
 func (c *Multi) Unwatch(keys ...string) *StatusCmd {
 	args := make([]interface{}, 1+len(keys))
 	args[0] = "UNWATCH"
@@ -66,6 +76,7 @@ func (c *Multi) Unwatch(keys ...string) *StatusCmd {
 	return cmd
 }
 
+// Discard discards queued commands.
 func (c *Multi) Discard() error {
 	if c.cmds == nil {
 		return errDiscard
@@ -74,10 +85,20 @@ func (c *Multi) Discard() error {
 	return nil
 }
 
+// Exec executes all previously queued commands in a transaction
+// and restores the connection state to normal.
+//
+// When using WATCH, EXEC will execute commands only if the watched keys
+// were not modified, allowing for a check-and-set mechanism.
+//
 // Exec always returns list of commands. If transaction fails
 // TxFailedErr is returned. Otherwise Exec returns error of the first
 // failed command or nil.
 func (c *Multi) Exec(f func() error) ([]Cmder, error) {
+	if c.closed {
+		return nil, errClosed
+	}
+
 	c.cmds = []Cmder{NewStatusCmd("MULTI")}
 	if err := f(); err != nil {
 		return nil, err
diff --git a/pipeline.go b/pipeline.go
index d7d13042..27c9f8ef 100644
--- a/pipeline.go
+++ b/pipeline.go
@@ -1,15 +1,22 @@
 package redis
 
+import (
+	"sync"
+	"sync/atomic"
+)
+
 // Pipeline implements pipelining as described in
-// http://redis.io/topics/pipelining. It's NOT safe for concurrent use
+// http://redis.io/topics/pipelining. It's safe for concurrent use
 // by multiple goroutines.
 type Pipeline struct {
 	commandable
 
 	client *baseClient
 
-	cmds   []Cmder
-	closed bool
+	mu   sync.Mutex // protects cmds
+	cmds []Cmder
+
+	closed int32
 }
 
 func (c *Client) Pipeline() *Pipeline {
@@ -27,36 +34,51 @@ func (c *Client) Pipelined(fn func(*Pipeline) error) ([]Cmder, error) {
 		return nil, err
 	}
 	cmds, err := pipe.Exec()
-	pipe.Close()
+	_ = pipe.Close()
 	return cmds, err
 }
 
 func (pipe *Pipeline) process(cmd Cmder) {
+	pipe.mu.Lock()
 	pipe.cmds = append(pipe.cmds, cmd)
+	pipe.mu.Unlock()
 }
 
 // Close closes the pipeline, releasing any open resources.
 func (pipe *Pipeline) Close() error {
+	atomic.StoreInt32(&pipe.closed, 1)
 	pipe.Discard()
-	pipe.closed = true
 	return nil
 }
 
+func (pipe *Pipeline) isClosed() bool {
+	return atomic.LoadInt32(&pipe.closed) == 1
+}
+
 // Discard resets the pipeline and discards queued commands.
 func (pipe *Pipeline) Discard() error {
-	if pipe.closed {
+	defer pipe.mu.Unlock()
+	pipe.mu.Lock()
+	if pipe.isClosed() {
 		return errClosed
 	}
 	pipe.cmds = pipe.cmds[:0]
 	return nil
 }
 
+// Exec executes all previously queued commands using one
+// client-server roundtrip.
+//
 // Exec always returns list of commands and error of the first failed
 // command if any.
 func (pipe *Pipeline) Exec() (cmds []Cmder, retErr error) {
-	if pipe.closed {
+	if pipe.isClosed() {
 		return nil, errClosed
 	}
+
+	defer pipe.mu.Unlock()
+	pipe.mu.Lock()
+
 	if len(pipe.cmds) == 0 {
 		return pipe.cmds, nil
 	}
diff --git a/pipeline_test.go b/pipeline_test.go
index ddf7480b..ed01baf4 100644
--- a/pipeline_test.go
+++ b/pipeline_test.go
@@ -150,4 +150,25 @@ var _ = Describe("Pipelining", func() {
 		wg.Wait()
 	})
 
+	It("should be thread-safe", func() {
+		const N = 1000
+
+		pipeline := client.Pipeline()
+		wg := &sync.WaitGroup{}
+		wg.Add(N)
+		for i := 0; i < N; i++ {
+			go func() {
+				pipeline.Ping()
+				wg.Done()
+			}()
+		}
+		wg.Wait()
+
+		cmds, err := pipeline.Exec()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(cmds).To(HaveLen(N))
+
+		Expect(pipeline.Close()).NotTo(HaveOccurred())
+	})
+
 })