From ef82e3705c2b3c748535267d87bc9ea16e621bec Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Wed, 10 Jun 2020 10:36:22 +0300
Subject: [PATCH] Add OnConnect context

---
 cluster.go    |  2 +-
 options.go    |  2 +-
 redis.go      |  2 +-
 redis_test.go |  2 +-
 ring.go       |  5 ++++-
 sentinel.go   |  5 +++--
 universal.go  | 27 ++++++++++++++++-----------
 7 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/cluster.go b/cluster.go
index aea830e..6f25c0c 100644
--- a/cluster.go
+++ b/cluster.go
@@ -55,7 +55,7 @@ type ClusterOptions struct {
 
 	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
 
-	OnConnect func(*Conn) error
+	OnConnect func(ctx context.Context, cn *Conn) error
 
 	Username string
 	Password string
diff --git a/options.go b/options.go
index 8aea428..1feb263 100644
--- a/options.go
+++ b/options.go
@@ -39,7 +39,7 @@ type Options struct {
 	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
 
 	// Hook that is called when new connection is established.
-	OnConnect func(*Conn) error
+	OnConnect func(ctx context.Context, cn *Conn) error
 
 	// Use the specified Username to authenticate the current connection with one of the connections defined in the ACL
 	// list when connecting to a Redis 6.0 instance, or greater, that is using the Redis ACL system.
diff --git a/redis.go b/redis.go
index 66e4c6a..569188b 100644
--- a/redis.go
+++ b/redis.go
@@ -269,7 +269,7 @@ func (c *baseClient) initConn(ctx context.Context, cn *pool.Conn) error {
 	}
 
 	if c.opt.OnConnect != nil {
-		return c.opt.OnConnect(conn)
+		return c.opt.OnConnect(ctx, conn)
 	}
 	return nil
 }
diff --git a/redis_test.go b/redis_test.go
index d5a3885..27ce351 100644
--- a/redis_test.go
+++ b/redis_test.go
@@ -372,7 +372,7 @@ var _ = Describe("Client OnConnect", func() {
 	BeforeEach(func() {
 		opt := redisOptions()
 		opt.DB = 0
-		opt.OnConnect = func(cn *redis.Conn) error {
+		opt.OnConnect = func(ctx context.Context, cn *redis.Conn) error {
 			return cn.ClientSetName(ctx, "on_connect").Err()
 		}
 
diff --git a/ring.go b/ring.go
index 9e02f94..99a3466 100644
--- a/ring.go
+++ b/ring.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/rand"
+	"net"
 	"strconv"
 	"sync"
 	"sync/atomic"
@@ -62,7 +63,8 @@ type RingOptions struct {
 
 	// Following options are copied from Options struct.
 
-	OnConnect func(*Conn) error
+	Dialer    func(ctx context.Context, network, addr string) (net.Conn, error)
+	OnConnect func(ctx context.Context, cn *Conn) error
 
 	Username string
 	DB       int
@@ -115,6 +117,7 @@ func (opt *RingOptions) init() {
 
 func (opt *RingOptions) clientOptions() *Options {
 	return &Options{
+		Dialer:    opt.Dialer,
 		OnConnect: opt.OnConnect,
 
 		DB: opt.DB,
diff --git a/sentinel.go b/sentinel.go
index d80f1e8..a6dbbae 100644
--- a/sentinel.go
+++ b/sentinel.go
@@ -28,7 +28,7 @@ type FailoverOptions struct {
 	// Following options are copied from Options struct.
 
 	Dialer    func(ctx context.Context, network, addr string) (net.Conn, error)
-	OnConnect func(*Conn) error
+	OnConnect func(ctx context.Context, cn *Conn) error
 
 	Username string
 	Password string
@@ -54,7 +54,8 @@ type FailoverOptions struct {
 
 func (opt *FailoverOptions) options() *Options {
 	return &Options{
-		Addr:      "FailoverClient",
+		Addr: "FailoverClient",
+
 		Dialer:    opt.Dialer,
 		OnConnect: opt.OnConnect,
 
diff --git a/universal.go b/universal.go
index 2107a81..6288d75 100644
--- a/universal.go
+++ b/universal.go
@@ -20,23 +20,28 @@ type UniversalOptions struct {
 
 	// Common options.
 
-	Dialer             func(ctx context.Context, network, addr string) (net.Conn, error)
-	OnConnect          func(*Conn) error
-	Username           string
-	Password           string
-	MaxRetries         int
-	MinRetryBackoff    time.Duration
-	MaxRetryBackoff    time.Duration
-	DialTimeout        time.Duration
-	ReadTimeout        time.Duration
-	WriteTimeout       time.Duration
+	Dialer    func(ctx context.Context, network, addr string) (net.Conn, error)
+	OnConnect func(ctx context.Context, cn *Conn) error
+
+	Username string
+	Password string
+
+	MaxRetries      int
+	MinRetryBackoff time.Duration
+	MaxRetryBackoff time.Duration
+
+	DialTimeout  time.Duration
+	ReadTimeout  time.Duration
+	WriteTimeout time.Duration
+
 	PoolSize           int
 	MinIdleConns       int
 	MaxConnAge         time.Duration
 	PoolTimeout        time.Duration
 	IdleTimeout        time.Duration
 	IdleCheckFrequency time.Duration
-	TLSConfig          *tls.Config
+
+	TLSConfig *tls.Config
 
 	// Only cluster clients.