Repository: gomodule/redigo
Pull request #470 (2020-04-26 13:09:41)

Handle `ctx` cancellation promptly (when waiting for a vacant connection from the pool)
I've also stumbled on [this one](https://github.com/gomodule/redigo/pull/426) on a project I've been working on recently. It might lead to pretty nasty scenarios when service under load uses `redigo`.

For some context you can check out [this reddit tread](https://www.reddit.com/r/golang/comments/4ibs68/why_is_the_select_statement_nondeterministic/).

LasTshaMAN: I also should note the following:

_if you are using `redis.Pool.Dial` and not `redis.Pool.DialContext` (that handles `ctx` cancellation  properly) your `go-routine` might still get stuck when dialing for a connection (even if the context you are using when calling `redis.Pool.GetContext` is cancelled). If your `redis.Pool.Dial` doesn't implement timeouts your `go-routine` will stuck forever. So it is something to think about._ 

Any reasons on why `redigo` doesn't provide default `ctx`-based `redis.Pool.DialContext` implementation similar to https://github.com/gomodule/redigo/blob/master/redis/conn.go#L179-L251 ? If there is no inherent problems with it, I could drop a PR ...
stevenh: Hi @LasTshaMAN could you rebase is to fix the conflicts. Will re-review once that's done.
LasTshaMAN: Sure, rebased.

@stevenh please check out this PR once again.

Patch:
diff --git a/redis/conn.go b/redis/conn.go
index a66e757..33b43be 100644
--- a/redis/conn.go
+++ b/redis/conn.go
@@ -80,6 +80,7 @@ type dialOptions struct {
 	dialer       *net.Dialer
 	dialContext  func(ctx context.Context, network, addr string) (net.Conn, error)
 	db           int
+	username     string
 	password     string
 	clientName   string
 	useTLS       bool
@@ -154,6 +155,14 @@ func DialPassword(password string) DialOption {
 	}}
 }
 
+// DialUsername specifies the username to use when connecting to
+// the Redis server when Redis ACLs are used.
+func DialUsername(username string) DialOption {
+	return DialOption{func(do *dialOptions) {
+		do.username = username
+	}}
+}
+
 // DialClientName specifies a client name to be used
 // by the Redis server connection.
 func DialClientName(name string) DialOption {
@@ -245,7 +254,12 @@ func DialContext(ctx context.Context, network, address string, options ...DialOp
 	}
 
 	if do.password != "" {
-		if _, err := c.Do("AUTH", do.password); err != nil {
+		authArgs := make([]interface{}, 0, 2)
+		if do.username != "" {
+			authArgs = append(authArgs, do.username)
+		}
+		authArgs = append(authArgs, do.password)
+		if _, err := c.Do("AUTH", authArgs...); err != nil {
 			netConn.Close()
 			return nil, err
 		}
@@ -303,7 +317,7 @@ func DialURL(rawurl string, options ...DialOption) (Conn, error) {
 	if u.User != nil {
 		password, isSet := u.User.Password()
 		if isSet {
-			options = append(options, DialPassword(password))
+			options = append(options, DialUsername(u.User.Username()), DialPassword(password))
 		}
 	}
 
diff --git a/redis/pool.go b/redis/pool.go
index 254c034..fef6bae 100644
--- a/redis/pool.go
+++ b/redis/pool.go
@@ -199,27 +199,10 @@ func (p *Pool) Get() Conn {
 // If the function completes without error, then the application must close the
 // returned connection.
 func (p *Pool) GetContext(ctx context.Context) (Conn, error) {
-	// Handle limit for p.Wait == true.
-	var waited time.Duration
-	if p.Wait && p.MaxActive > 0 {
-		p.lazyInit()
-
-		// wait indicates if we believe it will block so its not 100% accurate
-		// however for stats it should be good enough.
-		wait := len(p.ch) == 0
-		var start time.Time
-		if wait {
-			start = time.Now()
-		}
-		select {
-		case <-p.ch:
-		case <-ctx.Done():
-			err := ctx.Err()
-			return errorConn{err}, err
-		}
-		if wait {
-			waited = time.Since(start)
-		}
+	// Wait until there is a vacant connection in the pool.
+	waited, err := p.waitVacantConn(ctx)
+	if err != nil {
+		return nil, err
 	}
 
 	p.mu.Lock()
@@ -376,6 +359,51 @@ func (p *Pool) lazyInit() {
 	p.mu.Unlock()
 }
 
+// waitVacantConn waits for a vacant connection in pool if waiting
+// is enabled and pool size is limited, otherwise returns instantly.
+// If ctx expires before that, an error is returned.
+//
+// If there were no vacant connection in the pool right away it returns the time spent waiting
+// for that connection to appear in the pool.
+func (p *Pool) waitVacantConn(ctx context.Context) (waited time.Duration, err error) {
+	if !p.Wait || p.MaxActive <= 0 {
+		// No wait or no connection limit.
+		return 0, nil
+	}
+
+	p.lazyInit()
+
+	// wait indicates if we believe it will block so its not 100% accurate
+	// however for stats it should be good enough.
+	wait := len(p.ch) == 0
+	var start time.Time
+	if wait {
+		start = time.Now()
+	}
+
+	if ctx == nil {
+		<-p.ch
+	} else {
+		select {
+		case <-p.ch:
+			// Additionally check that context hasn't expired while we were waiting,
+			// because `select` picks a random `case` if several of them are "ready".
+			select {
+			case <-ctx.Done():
+				return 0, ctx.Err()
+			default:
+			}
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		}
+	}
+
+	if wait {
+		return time.Since(start), nil
+	}
+	return 0, nil
+}
+
 func (p *Pool) dial(ctx context.Context) (Conn, error) {
 	if p.DialContext != nil {
 		return p.DialContext(ctx)


Test patch:
diff --git a/redis/conn_test.go b/redis/conn_test.go
index 5864aca..dbc66e7 100644
--- a/redis/conn_test.go
+++ b/redis/conn_test.go
@@ -636,7 +636,9 @@ var dialURLTests = []struct {
 	r           string
 	w           string
 }{
-	{"password", "redis://x:abc123@localhost", "+OK\r\n", "*2\r\n$4\r\nAUTH\r\n$6\r\nabc123\r\n"},
+	{"password", "redis://:abc123@localhost", "+OK\r\n", "*2\r\n$4\r\nAUTH\r\n$6\r\nabc123\r\n"},
+	{"username and password", "redis://user:password@localhost", "+OK\r\n", "*3\r\n$4\r\nAUTH\r\n$4\r\nuser\r\n$8\r\npassword\r\n"},
+	{"username", "redis://x:@localhost", "+OK\r\n", ""},
 	{"database 3", "redis://localhost/3", "+OK\r\n", "*2\r\n$6\r\nSELECT\r\n$1\r\n3\r\n"},
 	{"database 99", "redis://localhost/99", "+OK\r\n", "*2\r\n$6\r\nSELECT\r\n$2\r\n99\r\n"},
 	{"no database", "redis://localhost/", "+OK\r\n", ""},
@@ -711,6 +713,36 @@ func TestDialTLSSKipVerify(t *testing.T) {
 	checkPingPong(t, &buf, c)
 }
 
+func TestDialUseACL(t *testing.T) {
+	var buf bytes.Buffer
+	_, err := redis.Dial("tcp", "localhost:6379",
+		redis.DialUsername("user"),
+		redis.DialPassword("password"),
+		dialTestConn(pingResponse, &buf))
+	if err != nil {
+		t.Fatal("dial error:", err)
+	}
+	if err != nil {
+		t.Fatal("dial error:", err)
+	}
+	expected := "*3\r\n$4\r\nAUTH\r\n$4\r\nuser\r\n$8\r\npassword\r\n"
+	if w := buf.String(); w != expected {
+		t.Errorf("got %q, want %q", w, expected)
+	}
+}
+
+// Connect to an Redis instance using the Redis ACL system
+func ExampleDial_acl() {
+	c, err := redis.Dial("tcp", "localhost:6379",
+		redis.DialUsername("username"),
+		redis.DialPassword("password"),
+	)
+	if err != nil {
+		// handle error
+	}
+	defer c.Close()
+}
+
 func TestDialClientName(t *testing.T) {
 	var buf bytes.Buffer
 	_, err := redis.Dial("tcp", ":6379",
diff --git a/redis/pool_test.go b/redis/pool_test.go
index 3e46bc4..a8efa51 100644
--- a/redis/pool_test.go
+++ b/redis/pool_test.go
@@ -899,20 +899,38 @@ func TestWaitPoolGetAfterClose(t *testing.T) {
 }
 
 func TestWaitPoolGetCanceledContext(t *testing.T) {
-	d := poolDialer{t: t}
-	p := &redis.Pool{
-		MaxIdle:   1,
-		MaxActive: 1,
-		Dial:      d.dial,
-		Wait:      true,
-	}
-	defer p.Close()
-	ctx, f := context.WithCancel(context.Background())
-	f()
-	c := p.Get()
-	defer c.Close()
-	_, err := p.GetContext(ctx)
-	if err != context.Canceled {
-		t.Fatalf("got error %v, want %v", err, context.Canceled)
-	}
+	t.Run("without vacant connection in the pool", func(t *testing.T) {
+		d := poolDialer{t: t}
+		p := &redis.Pool{
+			MaxIdle:   1,
+			MaxActive: 1,
+			Dial:      d.dial,
+			Wait:      true,
+		}
+		defer p.Close()
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		c := p.Get()
+		defer c.Close()
+		_, err := p.GetContext(ctx)
+		if err != context.Canceled {
+			t.Fatalf("got error %v, want %v", err, context.Canceled)
+		}
+	})
+	t.Run("with vacant connection in the pool", func(t *testing.T) {
+		d := poolDialer{t: t}
+		p := &redis.Pool{
+			MaxIdle:   1,
+			MaxActive: 1,
+			Dial:      d.dial,
+			Wait:      true,
+		}
+		defer p.Close()
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := p.GetContext(ctx)
+		if err != context.Canceled {
+			t.Fatalf("got error %v, want %v", err, context.Canceled)
+		}
+	})
 }