Repository: gomodule/redigo (go)

Issue:
Add DialContext function
I want Context to be used when dialing.
So I want DialContext.

By default you should use Dialer.DialContext.
https://golang.org/pkg/net/#Dialer.DialContext

Pull request:
Add DialContext function
Fixes #475.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 0884860..901ce78 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,5 +1,4 @@
 language: go
-sudo: false
 services:
   - redis-server
 
@@ -8,6 +7,8 @@ go:
   - 1.10.x
   - 1.11.x
   - 1.12.x
+  - 1.13.x
+  - 1.14.x
   - tip
 
 matrix:
diff --git a/redis/conn.go b/redis/conn.go
index 962c295..a66e757 100644
--- a/redis/conn.go
+++ b/redis/conn.go
@@ -17,6 +17,7 @@ package redis
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -77,7 +78,7 @@ type dialOptions struct {
 	readTimeout  time.Duration
 	writeTimeout time.Duration
 	dialer       *net.Dialer
-	dial         func(network, addr string) (net.Conn, error)
+	dialContext  func(ctx context.Context, network, addr string) (net.Conn, error)
 	db           int
 	password     string
 	clientName   string
@@ -123,7 +124,18 @@ func DialKeepAlive(d time.Duration) DialOption {
 // DialNetDial overrides DialConnectTimeout and DialKeepAlive.
 func DialNetDial(dial func(network, addr string) (net.Conn, error)) DialOption {
 	return DialOption{func(do *dialOptions) {
-		do.dial = dial
+		do.dialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
+			return dial(network, addr)
+		}
+	}}
+}
+
+// DialContextFunc specifies a custom dial function with context for creating TCP
+// connections, otherwise a net.Dialer customized via the other options is used.
+// DialContextFunc overrides DialConnectTimeout and DialKeepAlive.
+func DialContextFunc(f func(ctx context.Context, network, addr string) (net.Conn, error)) DialOption {
+	return DialOption{func(do *dialOptions) {
+		do.dialContext = f
 	}}
 }
 
@@ -177,6 +189,12 @@ func DialUseTLS(useTLS bool) DialOption {
 // Dial connects to the Redis server at the given network and
 // address using the specified options.
 func Dial(network, address string, options ...DialOption) (Conn, error) {
+	return DialContext(context.Background(), network, address, options...)
+}
+
+// DialContext connects to the Redis server at the given network and
+// address using the specified options and context.
+func DialContext(ctx context.Context, network, address string, options ...DialOption) (Conn, error) {
 	do := dialOptions{
 		dialer: &net.Dialer{
 			KeepAlive: time.Minute * 5,
@@ -185,11 +203,11 @@ func Dial(network, address string, options ...DialOption) (Conn, error) {
 	for _, option := range options {
 		option.f(&do)
 	}
-	if do.dial == nil {
-		do.dial = do.dialer.Dial
+	if do.dialContext == nil {
+		do.dialContext = do.dialer.DialContext
 	}
 
-	netConn, err := do.dial(network, address)
+	netConn, err := do.dialContext(ctx, network, address)
 	if err != nil {
 		return nil, err
 	}
diff --git a/redis/pool.go b/redis/pool.go
index cee0352..254c034 100644
--- a/redis/pool.go
+++ b/redis/pool.go
@@ -185,11 +185,9 @@ func NewPool(newFn func() (Conn, error), maxIdle int) *Pool {
 // getting an underlying connection, then the connection Err, Do, Send, Flush
 // and Receive methods return that error.
 func (p *Pool) Get() Conn {
-	pc, err := p.get(nil)
-	if err != nil {
-		return errorConn{err}
-	}
-	return &activeConn{p: p, pc: pc}
+	// GetContext returns errorConn in the first argument when an error occurs.
+	c, _ := p.GetContext(context.Background())
+	return c
 }
 
 // GetContext gets a connection using the provided context.
@@ -201,11 +199,90 @@ 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) {
-	pc, err := p.get(ctx)
+	// 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)
+		}
+	}
+
+	p.mu.Lock()
+
+	if waited > 0 {
+		p.waitCount++
+		p.waitDuration += waited
+	}
+
+	// Prune stale connections at the back of the idle list.
+	if p.IdleTimeout > 0 {
+		n := p.idle.count
+		for i := 0; i < n && p.idle.back != nil && p.idle.back.t.Add(p.IdleTimeout).Before(nowFunc()); i++ {
+			pc := p.idle.back
+			p.idle.popBack()
+			p.mu.Unlock()
+			pc.c.Close()
+			p.mu.Lock()
+			p.active--
+		}
+	}
+
+	// Get idle connection from the front of idle list.
+	for p.idle.front != nil {
+		pc := p.idle.front
+		p.idle.popFront()
+		p.mu.Unlock()
+		if (p.TestOnBorrow == nil || p.TestOnBorrow(pc.c, pc.t) == nil) &&
+			(p.MaxConnLifetime == 0 || nowFunc().Sub(pc.created) < p.MaxConnLifetime) {
+			return &activeConn{p: p, pc: pc}, nil
+		}
+		pc.c.Close()
+		p.mu.Lock()
+		p.active--
+	}
+
+	// Check for pool closed before dialing a new connection.
+	if p.closed {
+		p.mu.Unlock()
+		err := errors.New("redigo: get on closed pool")
+		return errorConn{err}, err
+	}
+
+	// Handle limit for p.Wait == false.
+	if !p.Wait && p.MaxActive > 0 && p.active >= p.MaxActive {
+		p.mu.Unlock()
+		return errorConn{ErrPoolExhausted}, ErrPoolExhausted
+	}
+
+	p.active++
+	p.mu.Unlock()
+	c, err := p.dial(ctx)
 	if err != nil {
+		c = nil
+		p.mu.Lock()
+		p.active--
+		if p.ch != nil && !p.closed {
+			p.ch <- struct{}{}
+		}
+		p.mu.Unlock()
 		return errorConn{err}, err
 	}
-	return &activeConn{p: p, pc: pc}, nil
+	return &activeConn{p: p, pc: &poolConn{c: c, created: nowFunc()}}, nil
 }
 
 // PoolStats contains pool statistics.
@@ -299,97 +376,6 @@ func (p *Pool) lazyInit() {
 	p.mu.Unlock()
 }
 
-// get prunes stale connections and returns a connection from the idle list or
-// creates a new connection.
-func (p *Pool) get(ctx context.Context) (*poolConn, 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()
-		}
-		if ctx == nil {
-			<-p.ch
-		} else {
-			select {
-			case <-p.ch:
-			case <-ctx.Done():
-				return nil, ctx.Err()
-			}
-		}
-		if wait {
-			waited = time.Since(start)
-		}
-	}
-
-	p.mu.Lock()
-
-	if waited > 0 {
-		p.waitCount++
-		p.waitDuration += waited
-	}
-
-	// Prune stale connections at the back of the idle list.
-	if p.IdleTimeout > 0 {
-		n := p.idle.count
-		for i := 0; i < n && p.idle.back != nil && p.idle.back.t.Add(p.IdleTimeout).Before(nowFunc()); i++ {
-			pc := p.idle.back
-			p.idle.popBack()
-			p.mu.Unlock()
-			pc.c.Close()
-			p.mu.Lock()
-			p.active--
-		}
-	}
-
-	// Get idle connection from the front of idle list.
-	for p.idle.front != nil {
-		pc := p.idle.front
-		p.idle.popFront()
-		p.mu.Unlock()
-		if (p.TestOnBorrow == nil || p.TestOnBorrow(pc.c, pc.t) == nil) &&
-			(p.MaxConnLifetime == 0 || nowFunc().Sub(pc.created) < p.MaxConnLifetime) {
-			return pc, nil
-		}
-		pc.c.Close()
-		p.mu.Lock()
-		p.active--
-	}
-
-	// Check for pool closed before dialing a new connection.
-	if p.closed {
-		p.mu.Unlock()
-		return nil, errors.New("redigo: get on closed pool")
-	}
-
-	// Handle limit for p.Wait == false.
-	if !p.Wait && p.MaxActive > 0 && p.active >= p.MaxActive {
-		p.mu.Unlock()
-		return nil, ErrPoolExhausted
-	}
-
-	p.active++
-	p.mu.Unlock()
-	c, err := p.dial(ctx)
-	if err != nil {
-		c = nil
-		p.mu.Lock()
-		p.active--
-		if p.ch != nil && !p.closed {
-			p.ch <- struct{}{}
-		}
-		p.mu.Unlock()
-	}
-	return &poolConn{c: c, created: nowFunc()}, err
-}
-
 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 08da5fa..5864aca 100644
--- a/redis/conn_test.go
+++ b/redis/conn_test.go
@@ -16,6 +16,7 @@ package redis_test
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -531,6 +532,37 @@ func TestReadTimeout(t *testing.T) {
 	}
 }
 
+func TestDialContextFunc(t *testing.T) {
+	var isPassed bool
+	f := func(ctx context.Context, network, addr string) (net.Conn, error) {
+		isPassed = true
+		return &testConn{}, nil
+	}
+
+	_, err := redis.DialContext(context.Background(), "", "", redis.DialContextFunc(f))
+	if err != nil {
+		t.Fatalf("DialContext returned %v", err)
+	}
+
+	if !isPassed {
+		t.Fatal("DialContextFunc not passed")
+	}
+}
+
+func TestDialContext_CanceledContext(t *testing.T) {
+	addr, err := redis.DefaultServerAddr()
+	if err != nil {
+		t.Fatalf("redis.DefaultServerAddr returned %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	if _, err = redis.DialContext(ctx, "tcp", addr); err == nil {
+		t.Fatalf("DialContext returned nil, expect error")
+	}
+}
+
 var dialErrors = []struct {
 	rawurl        string
 	expectedError string
@@ -725,6 +757,16 @@ func ExampleDial() {
 	defer c.Close()
 }
 
+// Connect to local instance of Redis running on the default port using the provided context.
+func ExampleDialContext() {
+	ctx := context.Background()
+	c, err := redis.DialContext(ctx, "tcp", ":6379")
+	if err != nil {
+		// handle error
+	}
+	defer c.Close()
+}
+
 // Connect to remote instance of Redis using a URL.
 func ExampleDialURL() {
 	c, err := redis.DialURL(os.Getenv("REDIS_URL"))
diff --git a/redis/pool_test.go b/redis/pool_test.go
index 3864d2a..3e46bc4 100644
--- a/redis/pool_test.go
+++ b/redis/pool_test.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net"
 	"reflect"
 	"sync"
 	"testing"
@@ -840,6 +841,48 @@ func TestWaitPoolGetContextWithDialContext(t *testing.T) {
 	defer c.Close()
 }
 
+func TestPoolGetContext_DialContext(t *testing.T) {
+	var isPassed bool
+	f := func(ctx context.Context, network, addr string) (net.Conn, error) {
+		isPassed = true
+		return &testConn{}, nil
+	}
+
+	p := &redis.Pool{
+		DialContext: func(ctx context.Context) (redis.Conn, error) {
+			return redis.DialContext(ctx, "", "", redis.DialContextFunc(f))
+		},
+	}
+	defer p.Close()
+
+	if _, err := p.GetContext(context.Background()); err != nil {
+		t.Fatalf("GetContext returned %v", err)
+	}
+
+	if !isPassed {
+		t.Fatal("DialContextFunc not passed")
+	}
+}
+
+func TestPoolGetContext_DialContext_CanceledContext(t *testing.T) {
+	addr, err := redis.DefaultServerAddr()
+	if err != nil {
+		t.Fatalf("redis.DefaultServerAddr returned %v", err)
+	}
+
+	p := &redis.Pool{
+		DialContext: func(ctx context.Context) (redis.Conn, error) { return redis.DialContext(ctx, "tcp", addr) },
+	}
+	defer p.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	if _, err := p.GetContext(ctx); err == nil {
+		t.Fatalf("GetContext returned nil, expect error")
+	}
+}
+
 func TestWaitPoolGetAfterClose(t *testing.T) {
 	d := poolDialer{t: t}
 	p := &redis.Pool{