Make Pipeline thread-safe. Fixes #166.

This commit is contained in:
Vladimir Mihailenco 2015-11-04 14:25:48 +02:00
parent a2dba7df0c
commit d0d3920e69
3 changed files with 74 additions and 10 deletions

View File

@ -9,13 +9,18 @@ import (
var errDiscard = errors.New("redis: Discard can be used only inside Exec") var errDiscard = errors.New("redis: Discard can be used only inside Exec")
// Multi implements Redis transactions as described in // Multi implements Redis transactions as described in
// http://redis.io/topics/transactions. It's NOT safe for concurrent // http://redis.io/topics/transactions. It's NOT safe for concurrent use
// use by multiple goroutines. // 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 { type Multi struct {
commandable commandable
base *baseClient base *baseClient
cmds []Cmder
cmds []Cmder
closed bool
} }
func (c *Client) Multi() *Multi { 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 { func (c *Multi) Close() error {
c.closed = true
if err := c.Unwatch().Err(); err != nil { if err := c.Unwatch().Err(); err != nil {
log.Printf("redis: Unwatch failed: %s", err) log.Printf("redis: Unwatch failed: %s", err)
} }
return c.base.Close() return c.base.Close()
} }
// Watch marks the keys to be watched for conditional execution
// of a transaction.
func (c *Multi) Watch(keys ...string) *StatusCmd { func (c *Multi) Watch(keys ...string) *StatusCmd {
args := make([]interface{}, 1+len(keys)) args := make([]interface{}, 1+len(keys))
args[0] = "WATCH" args[0] = "WATCH"
@ -55,6 +64,7 @@ func (c *Multi) Watch(keys ...string) *StatusCmd {
return cmd return cmd
} }
// Unwatch flushes all the previously watched keys for a transaction.
func (c *Multi) Unwatch(keys ...string) *StatusCmd { func (c *Multi) Unwatch(keys ...string) *StatusCmd {
args := make([]interface{}, 1+len(keys)) args := make([]interface{}, 1+len(keys))
args[0] = "UNWATCH" args[0] = "UNWATCH"
@ -66,6 +76,7 @@ func (c *Multi) Unwatch(keys ...string) *StatusCmd {
return cmd return cmd
} }
// Discard discards queued commands.
func (c *Multi) Discard() error { func (c *Multi) Discard() error {
if c.cmds == nil { if c.cmds == nil {
return errDiscard return errDiscard
@ -74,10 +85,20 @@ func (c *Multi) Discard() error {
return nil 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 // Exec always returns list of commands. If transaction fails
// TxFailedErr is returned. Otherwise Exec returns error of the first // TxFailedErr is returned. Otherwise Exec returns error of the first
// failed command or nil. // failed command or nil.
func (c *Multi) Exec(f func() error) ([]Cmder, error) { func (c *Multi) Exec(f func() error) ([]Cmder, error) {
if c.closed {
return nil, errClosed
}
c.cmds = []Cmder{NewStatusCmd("MULTI")} c.cmds = []Cmder{NewStatusCmd("MULTI")}
if err := f(); err != nil { if err := f(); err != nil {
return nil, err return nil, err

View File

@ -1,15 +1,22 @@
package redis package redis
import (
"sync"
"sync/atomic"
)
// Pipeline implements pipelining as described in // 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. // by multiple goroutines.
type Pipeline struct { type Pipeline struct {
commandable commandable
client *baseClient client *baseClient
cmds []Cmder mu sync.Mutex // protects cmds
closed bool cmds []Cmder
closed int32
} }
func (c *Client) Pipeline() *Pipeline { func (c *Client) Pipeline() *Pipeline {
@ -27,36 +34,51 @@ func (c *Client) Pipelined(fn func(*Pipeline) error) ([]Cmder, error) {
return nil, err return nil, err
} }
cmds, err := pipe.Exec() cmds, err := pipe.Exec()
pipe.Close() _ = pipe.Close()
return cmds, err return cmds, err
} }
func (pipe *Pipeline) process(cmd Cmder) { func (pipe *Pipeline) process(cmd Cmder) {
pipe.mu.Lock()
pipe.cmds = append(pipe.cmds, cmd) pipe.cmds = append(pipe.cmds, cmd)
pipe.mu.Unlock()
} }
// Close closes the pipeline, releasing any open resources. // Close closes the pipeline, releasing any open resources.
func (pipe *Pipeline) Close() error { func (pipe *Pipeline) Close() error {
atomic.StoreInt32(&pipe.closed, 1)
pipe.Discard() pipe.Discard()
pipe.closed = true
return nil return nil
} }
func (pipe *Pipeline) isClosed() bool {
return atomic.LoadInt32(&pipe.closed) == 1
}
// Discard resets the pipeline and discards queued commands. // Discard resets the pipeline and discards queued commands.
func (pipe *Pipeline) Discard() error { func (pipe *Pipeline) Discard() error {
if pipe.closed { defer pipe.mu.Unlock()
pipe.mu.Lock()
if pipe.isClosed() {
return errClosed return errClosed
} }
pipe.cmds = pipe.cmds[:0] pipe.cmds = pipe.cmds[:0]
return nil 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 // Exec always returns list of commands and error of the first failed
// command if any. // command if any.
func (pipe *Pipeline) Exec() (cmds []Cmder, retErr error) { func (pipe *Pipeline) Exec() (cmds []Cmder, retErr error) {
if pipe.closed { if pipe.isClosed() {
return nil, errClosed return nil, errClosed
} }
defer pipe.mu.Unlock()
pipe.mu.Lock()
if len(pipe.cmds) == 0 { if len(pipe.cmds) == 0 {
return pipe.cmds, nil return pipe.cmds, nil
} }

View File

@ -150,4 +150,25 @@ var _ = Describe("Pipelining", func() {
wg.Wait() 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())
})
}) })