Optimize Otel instrumentation

internal/pool/conn.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/go-redis/redis/v8/internal"
 	"github.com/go-redis/redis/v8/internal/proto"
-	"go.opentelemetry.io/otel/trace"
 )
 
 var noDeadline = time.Time{}
@@ -66,41 +65,43 @@ func (cn *Conn) RemoteAddr() net.Addr {
 }
 
 func (cn *Conn) WithReader(ctx context.Context, timeout time.Duration, fn func(rd *proto.Reader) error) error {
-	return internal.WithSpan(ctx, "redis.with_reader", func(ctx context.Context, span trace.Span) error {
+	ctx, span := internal.StartSpan(ctx, "redis.with_reader")
-		if err := cn.netConn.SetReadDeadline(cn.deadline(ctx, timeout)); err != nil {
+	defer span.End()
-			return internal.RecordError(ctx, span, err)
+
-		}
+	if err := cn.netConn.SetReadDeadline(cn.deadline(ctx, timeout)); err != nil {
-		if err := fn(cn.rd); err != nil {
+		return internal.RecordError(ctx, span, err)
-			return internal.RecordError(ctx, span, err)
+	}
-		}
+	if err := fn(cn.rd); err != nil {
-		return nil
+		return internal.RecordError(ctx, span, err)
-	})
+	}
+	return nil
 }
 
 func (cn *Conn) WithWriter(
 	ctx context.Context, timeout time.Duration, fn func(wr *proto.Writer) error,
 ) error {
-	return internal.WithSpan(ctx, "redis.with_writer", func(ctx context.Context, span trace.Span) error {
+	ctx, span := internal.StartSpan(ctx, "redis.with_writer")
-		if err := cn.netConn.SetWriteDeadline(cn.deadline(ctx, timeout)); err != nil {
+	defer span.End()
-			return internal.RecordError(ctx, span, err)
-		}
 
-		if cn.bw.Buffered() > 0 {
+	if err := cn.netConn.SetWriteDeadline(cn.deadline(ctx, timeout)); err != nil {
-			cn.bw.Reset(cn.netConn)
+		return internal.RecordError(ctx, span, err)
-		}
+	}
 
-		if err := fn(cn.wr); err != nil {
+	if cn.bw.Buffered() > 0 {
-			return internal.RecordError(ctx, span, err)
+		cn.bw.Reset(cn.netConn)
-		}
+	}
 
-		if err := cn.bw.Flush(); err != nil {
+	if err := fn(cn.wr); err != nil {
-			return internal.RecordError(ctx, span, err)
+		return internal.RecordError(ctx, span, err)
-		}
+	}
 
-		internal.WritesCounter.Add(ctx, 1)
+	if err := cn.bw.Flush(); err != nil {
+		return internal.RecordError(ctx, span, err)
+	}
 
-		return nil
+	internal.WritesCounter.Add(ctx, 1)
-	})
+
+	return nil
 }
 
 func (cn *Conn) Close() error {

internal/util.go

@@ -11,17 +11,18 @@ import (
 )
 
 func Sleep(ctx context.Context, dur time.Duration) error {
-	return WithSpan(ctx, "time.Sleep", func(ctx context.Context, span trace.Span) error {
+	_, span := StartSpan(ctx, "time.Sleep")
-		t := time.NewTimer(dur)
+	defer span.End()
-		defer t.Stop()
 
-		select {
+	t := time.NewTimer(dur)
-		case <-t.C:
+	defer t.Stop()
-			return nil
+
-		case <-ctx.Done():
+	select {
-			return ctx.Err()
+	case <-t.C:
-		}
+		return nil
-	})
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func ToLower(s string) string {
@@ -54,15 +55,11 @@ func isLower(s string) bool {
 
 var tracer = otel.Tracer("github.com/go-redis/redis")
 
-func WithSpan(ctx context.Context, name string, fn func(context.Context, trace.Span) error) error {
+func StartSpan(ctx context.Context, name string) (context.Context, trace.Span) {
 	if span := trace.SpanFromContext(ctx); !span.IsRecording() {
-		return fn(ctx, span)
+		return ctx, span
 	}
-
+	return tracer.Start(ctx, name)
-	ctx, span := tracer.Start(ctx, name)
-	defer span.End()
-
-	return fn(ctx, span)
 }
 
 func RecordError(ctx context.Context, span trace.Span, err error) error {

options.go

@@ -15,7 +15,6 @@ import (
 	"github.com/go-redis/redis/v8/internal"
 	"github.com/go-redis/redis/v8/internal/pool"
 	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/trace"
 )
 
 // Limiter is the interface of a rate limiter or a circuit breaker.
@@ -292,20 +291,21 @@ func getUserPassword(u *url.URL) (string, string) {
 func newConnPool(opt *Options) *pool.ConnPool {
 	return pool.NewConnPool(&pool.Options{
 		Dialer: func(ctx context.Context) (net.Conn, error) {
-			var conn net.Conn
+			ctx, span := internal.StartSpan(ctx, "redis.dial")
-			err := internal.WithSpan(ctx, "redis.dial", func(ctx context.Context, span trace.Span) error {
+			defer span.End()
+
+			if span.IsRecording() {
 				span.SetAttributes(
 					attribute.String("db.connection_string", opt.Addr),
 				)
+			}
 
-				var err error
+			cn, err := opt.Dialer(ctx, opt.Network, opt.Addr)
-				conn, err = opt.Dialer(ctx, opt.Network, opt.Addr)
+			if err != nil {
-				if err != nil {
+				return nil, internal.RecordError(ctx, span, err)
-					_ = internal.RecordError(ctx, span, err)
+			}
-				}
+
-				return err
+			return cn, nil
-			})
-			return conn, err
 		},
 		PoolSize:           opt.PoolSize,
 		MinIdleConns:       opt.MinIdleConns,

redis.go

@@ -11,7 +11,6 @@ import (
 	"github.com/go-redis/redis/v8/internal/pool"
 	"github.com/go-redis/redis/v8/internal/proto"
 	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/trace"
 )
 
 // Nil reply returned by Redis when key does not exist.
@@ -214,10 +213,7 @@ func (c *baseClient) _getConn(ctx context.Context) (*pool.Conn, error) {
 		return cn, nil
 	}
 
-	err = internal.WithSpan(ctx, "redis.init_conn", func(ctx context.Context, span trace.Span) error {
+	if err := c.initConn(ctx, cn); err != nil {
-		return c.initConn(ctx, cn)
-	})
-	if err != nil {
 		c.connPool.Remove(ctx, cn, err)
 		if err := errors.Unwrap(err); err != nil {
 			return nil, err
@@ -241,6 +237,9 @@ func (c *baseClient) initConn(ctx context.Context, cn *pool.Conn) error {
 		return nil
 	}
 
+	ctx, span := internal.StartSpan(ctx, "redis.init_conn")
+	defer span.End()
+
 	connPool := pool.NewSingleConnPool(c.connPool, cn)
 	conn := newConn(ctx, c.opt, connPool)
 
@@ -288,43 +287,44 @@ func (c *baseClient) releaseConn(ctx context.Context, cn *pool.Conn, err error)
 func (c *baseClient) withConn(
 	ctx context.Context, fn func(context.Context, *pool.Conn) error,
 ) error {
-	return internal.WithSpan(ctx, "redis.with_conn", func(ctx context.Context, span trace.Span) error {
+	ctx, span := internal.StartSpan(ctx, "redis.with_conn")
-		cn, err := c.getConn(ctx)
+	defer span.End()
-		if err != nil {
+
-			return err
+	cn, err := c.getConn(ctx)
+	if err != nil {
+		return err
+	}
+
+	if span.IsRecording() {
+		if remoteAddr := cn.RemoteAddr(); remoteAddr != nil {
+			span.SetAttributes(attribute.String("net.peer.ip", remoteAddr.String()))
 		}
+	}
 
-		if span.IsRecording() {
+	defer func() {
-			if remoteAddr := cn.RemoteAddr(); remoteAddr != nil {
+		c.releaseConn(ctx, cn, err)
-				span.SetAttributes(attribute.String("net.peer.ip", remoteAddr.String()))
+	}()
-			}
-		}
 
-		defer func() {
+	done := ctx.Done()
-			c.releaseConn(ctx, cn, err)
+	if done == nil {
-		}()
+		err = fn(ctx, cn)
+		return err
+	}
 
-		done := ctx.Done()
+	errc := make(chan error, 1)
-		if done == nil {
+	go func() { errc <- fn(ctx, cn) }()
-			err = fn(ctx, cn)
-			return err
-		}
 
-		errc := make(chan error, 1)
+	select {
-		go func() { errc <- fn(ctx, cn) }()
+	case <-done:
+		_ = cn.Close()
+		// Wait for the goroutine to finish and send something.
+		<-errc
 
-		select {
+		err = ctx.Err()
-		case <-done:
+		return err
-			_ = cn.Close()
+	case err = <-errc:
-			// Wait for the goroutine to finish and send something.
+		return err
-			<-errc
+	}
-
-			err = ctx.Err()
-			return err
-		case err = <-errc:
-			return err
-		}
-	})
 }
 
 func (c *baseClient) process(ctx context.Context, cmd Cmder) error {
@@ -332,47 +332,50 @@ func (c *baseClient) process(ctx context.Context, cmd Cmder) error {
 	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
 		attempt := attempt
 
-		var retry bool
+		retry, err := c._process(ctx, cmd, attempt)
-		err := internal.WithSpan(ctx, "redis.process", func(ctx context.Context, span trace.Span) error {
-			if attempt > 0 {
-				if err := internal.Sleep(ctx, c.retryBackoff(attempt)); err != nil {
-					return err
-				}
-			}
-
-			retryTimeout := uint32(1)
-			err := c.withConn(ctx, func(ctx context.Context, cn *pool.Conn) error {
-				err := cn.WithWriter(ctx, c.opt.WriteTimeout, func(wr *proto.Writer) error {
-					return writeCmd(wr, cmd)
-				})
-				if err != nil {
-					return err
-				}
-
-				err = cn.WithReader(ctx, c.cmdTimeout(cmd), cmd.readReply)
-				if err != nil {
-					if cmd.readTimeout() == nil {
-						atomic.StoreUint32(&retryTimeout, 1)
-					}
-					return err
-				}
-
-				return nil
-			})
-			if err == nil {
-				return nil
-			}
-			retry = shouldRetry(err, atomic.LoadUint32(&retryTimeout) == 1)
-			return err
-		})
 		if err == nil || !retry {
 			return err
 		}
+
 		lastErr = err
 	}
 	return lastErr
 }
 
+func (c *baseClient) _process(ctx context.Context, cmd Cmder, attempt int) (bool, error) {
+	if attempt > 0 {
+		if err := internal.Sleep(ctx, c.retryBackoff(attempt)); err != nil {
+			return false, err
+		}
+	}
+
+	retryTimeout := uint32(1)
+	err := c.withConn(ctx, func(ctx context.Context, cn *pool.Conn) error {
+		err := cn.WithWriter(ctx, c.opt.WriteTimeout, func(wr *proto.Writer) error {
+			return writeCmd(wr, cmd)
+		})
+		if err != nil {
+			return err
+		}
+
+		err = cn.WithReader(ctx, c.cmdTimeout(cmd), cmd.readReply)
+		if err != nil {
+			if cmd.readTimeout() == nil {
+				atomic.StoreUint32(&retryTimeout, 1)
+			}
+			return err
+		}
+
+		return nil
+	})
+	if err == nil {
+		return false, nil
+	}
+
+	retry := shouldRetry(err, atomic.LoadUint32(&retryTimeout) == 1)
+	return retry, err
+}
+
 func (c *baseClient) retryBackoff(attempt int) time.Duration {
 	return internal.RetryBackoff(attempt, c.opt.MinRetryBackoff, c.opt.MaxRetryBackoff)
 }