dial hook

Author: vmihailenco

example_instrumentation_test.go

@@ -3,6 +3,7 @@ package redis_test
 import (
 	"context"
 	"fmt"
+	"net"
 
 	"github.com/go-redis/redis/v9"
 )
@@ -12,8 +13,8 @@ type redisHook struct{}
 var _ redis.Hook = redisHook{}
 
 func (redisHook) DialHook(hook redis.DialHook) redis.DialHook {
-	return func(ctx context.Context) error {
-		return hook(ctx)
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		return hook(ctx, network, addr)
 	}
 }
 

extra/redisotel/metrics.go

@@ -2,6 +2,7 @@ package redisotel
 
 import (
 	"context"
+	"net"
 	"time"
 
 	"github.com/go-redis/redis/v9"
@@ -75,11 +76,11 @@ type metricsHook struct {
 var _ redis.Hook = (*metricsHook)(nil)
 
 func (mh *metricsHook) DialHook(hook redis.DialHook) redis.DialHook {
-	return func(ctx context.Context) error {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
 		start := time.Now()
-		err := hook(ctx)
+		conn, err := hook(ctx, network, addr)
 		mh.createTime.Record(ctx, time.Since(start).Microseconds())
-		return err
+		return conn, err
 	}
 }
 

extra/redisotel/tracing.go

@@ -2,6 +2,7 @@ package redisotel
 
 import (
 	"context"
+	"net"
 
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
@@ -41,19 +42,26 @@ func NewTracingHook(opts ...Option) *TracingHook {
 }
 
 func (th *TracingHook) DialHook(hook redis.DialHook) redis.DialHook {
-	return func(ctx context.Context) error {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
 		if !trace.SpanFromContext(ctx).IsRecording() {
-			return hook(ctx)
+			return hook(ctx, network, addr)
 		}
 
-		ctx, span := th.conf.tracer.Start(ctx, "redis.dial", th.spanOpts...)
+		spanOpts := th.spanOpts
+		spanOpts = append(spanOpts, trace.WithAttributes(
+			attribute.String("network", network),
+			attribute.String("addr", addr),
+		))
+
+		ctx, span := th.conf.tracer.Start(ctx, "redis.dial", spanOpts...)
 		defer span.End()
 
-		if err := hook(ctx); err != nil {
+		conn, err := hook(ctx, network, addr)
+		if err != nil {
 			recordError(ctx, span, err)
-			return err
+			return nil, err
 		}
-		return nil
+		return conn, nil
 	}
 }
 

options.go

@@ -470,10 +470,13 @@ func getUserPassword(u *url.URL) (string, string) {
 	return user, password
 }
 
-func newConnPool(opt *Options) *pool.ConnPool {
+func newConnPool(
+	opt *Options,
+	dialer func(ctx context.Context, network, addr string) (net.Conn, error),
+) *pool.ConnPool {
 	return pool.NewConnPool(&pool.Options{
 		Dialer: func(ctx context.Context) (net.Conn, error) {
-			return opt.Dialer(ctx, opt.Network, opt.Addr)
+			return dialer(ctx, opt.Network, opt.Addr)
 		},
 		PoolFIFO:        opt.PoolFIFO,
 		PoolSize:        opt.PoolSize,

redis.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -30,7 +31,7 @@ type Hook interface {
 }
 
 type (
-	DialHook            func(ctx context.Context) error
+	DialHook            func(ctx context.Context, network, addr string) (net.Conn, error)
 	ProcessHook         func(ctx context.Context, cmd Cmder) error
 	ProcessPipelineHook func(ctx context.Context, cmds []Cmder) error
 )
@@ -68,6 +69,15 @@ func (hs *hooks) clone() hooks {
 	return clone
 }
 
+func (hs *hooks) setDial(dial DialHook) {
+	hs.dial = dial
+	for _, h := range hs.slice {
+		if wrapped := h.DialHook(hs.dial); wrapped != nil {
+			hs.dial = wrapped
+		}
+	}
+}
+
 func (hs *hooks) setProcess(process ProcessHook) {
 	hs.process = process
 	for _, h := range hs.slice {
@@ -124,13 +134,6 @@ type baseClient struct {
 	onClose func() error // hook called when client is closed
 }
 
-func newBaseClient(opt *Options, connPool pool.Pooler) *baseClient {
-	return &baseClient{
-		opt:      opt,
-		connPool: connPool,
-	}
-}
-
 func (c *baseClient) clone() *baseClient {
 	clone := *c
 	return &clone
@@ -286,6 +289,10 @@ func (c *baseClient) withConn(
 	return fn(ctx, cn)
 }
 
+func (c *baseClient) dial(ctx context.Context, network, addr string) (net.Conn, error) {
+	return c.opt.Dialer(ctx, network, addr)
+}
+
 func (c *baseClient) process(ctx context.Context, cmd Cmder) error {
 	var lastErr error
 	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
@@ -527,15 +534,19 @@ func NewClient(opt *Options) *Client {
 	opt.init()
 
 	c := Client{
-		baseClient: newBaseClient(opt, newConnPool(opt)),
+		baseClient: &baseClient{
+			opt: opt,
+		},
 	}
+	c.connPool = newConnPool(opt, c.baseClient.dial)
 	c.init()
 
 	return &c
 }
 
 func (c *Client) init() {
 	c.cmdable = c.Process
+	c.hooks.setDial(c.baseClient.dial)
 	c.hooks.setProcess(c.baseClient.process)
 	c.hooks.setProcessPipeline(c.baseClient.processPipeline)
 	c.hooks.setProcessTxPipeline(c.baseClient.processTxPipeline)
@@ -696,6 +707,7 @@ func newConn(opt *Options, connPool pool.Pooler) *Conn {
 	c.cmdable = c.Process
 	c.statefulCmdable = c.Process
 
+	c.hooks.setDial(c.baseClient.dial)
 	c.hooks.setProcess(c.baseClient.process)
 	c.hooks.setProcessPipeline(c.baseClient.processPipeline)
 	c.hooks.setProcessTxPipeline(c.baseClient.processTxPipeline)

sentinel.go

@@ -198,7 +198,17 @@ func NewFailoverClient(failoverOpt *FailoverOptions) *Client {
 	opt.Dialer = masterReplicaDialer(failover)
 	opt.init()
 
-	connPool := newConnPool(opt)
+	var connPool *pool.ConnPool
+
+	rdb := &Client{
+		baseClient: &baseClient{
+			opt: opt,
+		},
+	}
+	connPool = newConnPool(opt, rdb.baseClient.dial)
+	rdb.connPool = connPool
+	rdb.onClose = failover.Close
+	rdb.init()
 
 	failover.mu.Lock()
 	failover.onFailover = func(ctx context.Context, addr string) {
@@ -208,12 +218,6 @@ func NewFailoverClient(failoverOpt *FailoverOptions) *Client {
 	}
 	failover.mu.Unlock()
 
-	rdb := &Client{
-		baseClient: newBaseClient(opt, connPool),
-	}
-	rdb.onClose = failover.Close
-	rdb.init()
-
 	return rdb
 }
 
@@ -262,11 +266,12 @@ func NewSentinelClient(opt *Options) *SentinelClient {
 	opt.init()
 	c := &SentinelClient{
 		baseClient: &baseClient{
-			opt:      opt,
-			connPool: newConnPool(opt),
+			opt: opt,
 		},
 	}
+	c.connPool = newConnPool(opt, c.baseClient.dial)
 
+	c.hooks.setDial(c.baseClient.dial)
 	c.hooks.setProcess(c.baseClient.process)
 
 	return c

tx.go

@@ -38,6 +38,7 @@ func (c *Tx) init() {
 	c.cmdable = c.Process
 	c.statefulCmdable = c.Process
 
+	c.hooks.setDial(c.baseClient.dial)
 	c.hooks.setProcess(c.baseClient.process)
 	c.hooks.setProcessPipeline(c.baseClient.processPipeline)
 	c.hooks.setProcessTxPipeline(c.baseClient.processTxPipeline)