From 566840fe273530afbab74d6c09d208d9ecbec5ab Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Mon, 28 Sep 2020 15:35:57 +0300
Subject: [PATCH] otel: add RemoteAddr

---
 internal/pool/conn.go | 10 +++++++---
 internal/util.go      | 10 +++++-----
 options.go            |  4 ++--
 redis.go              | 15 ++++++++++++---
 4 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/internal/pool/conn.go b/internal/pool/conn.go
index 9c243c6..b930a90 100644
--- a/internal/pool/conn.go
+++ b/internal/pool/conn.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/go-redis/redis/v8/internal"
 	"github.com/go-redis/redis/v8/internal/proto"
+	"go.opentelemetry.io/otel/api/trace"
 )
 
 var noDeadline = time.Time{}
@@ -58,11 +59,14 @@ func (cn *Conn) Write(b []byte) (int, error) {
 }
 
 func (cn *Conn) RemoteAddr() net.Addr {
-	return cn.netConn.RemoteAddr()
+	if cn.netConn != nil {
+		return cn.netConn.RemoteAddr()
+	}
+	return nil
 }
 
 func (cn *Conn) WithReader(ctx context.Context, timeout time.Duration, fn func(rd *proto.Reader) error) error {
-	return internal.WithSpan(ctx, "with_reader", func(ctx context.Context) error {
+	return internal.WithSpan(ctx, "with_reader", func(ctx context.Context, span trace.Span) error {
 		if err := cn.netConn.SetReadDeadline(cn.deadline(ctx, timeout)); err != nil {
 			return internal.RecordError(ctx, err)
 		}
@@ -76,7 +80,7 @@ func (cn *Conn) WithReader(ctx context.Context, timeout time.Duration, fn func(r
 func (cn *Conn) WithWriter(
 	ctx context.Context, timeout time.Duration, fn func(wr *proto.Writer) error,
 ) error {
-	return internal.WithSpan(ctx, "with_writer", func(ctx context.Context) error {
+	return internal.WithSpan(ctx, "with_writer", func(ctx context.Context, span trace.Span) error {
 		if err := cn.netConn.SetWriteDeadline(cn.deadline(ctx, timeout)); err != nil {
 			return internal.RecordError(ctx, err)
 		}
diff --git a/internal/util.go b/internal/util.go
index 710bfe1..9c861b0 100644
--- a/internal/util.go
+++ b/internal/util.go
@@ -11,7 +11,7 @@ import (
 )
 
 func Sleep(ctx context.Context, dur time.Duration) error {
-	return WithSpan(ctx, "sleep", func(ctx context.Context) error {
+	return WithSpan(ctx, "sleep", func(ctx context.Context, span trace.Span) error {
 		t := time.NewTimer(dur)
 		defer t.Stop()
 
@@ -62,15 +62,15 @@ func Unwrap(err error) error {
 
 //------------------------------------------------------------------------------
 
-func WithSpan(ctx context.Context, name string, fn func(context.Context) error) error {
-	if !trace.SpanFromContext(ctx).IsRecording() {
-		return fn(ctx)
+func WithSpan(ctx context.Context, name string, fn func(context.Context, trace.Span) error) error {
+	if span := trace.SpanFromContext(ctx); !span.IsRecording() {
+		return fn(ctx, span)
 	}
 
 	ctx, span := global.Tracer("github.com/go-redis/redis").Start(ctx, name)
 	defer span.End()
 
-	return fn(ctx)
+	return fn(ctx, span)
 }
 
 func RecordError(ctx context.Context, err error) error {
diff --git a/options.go b/options.go
index 32ff1ea..b3d3196 100644
--- a/options.go
+++ b/options.go
@@ -293,9 +293,9 @@ func newConnPool(opt *Options) *pool.ConnPool {
 	return pool.NewConnPool(&pool.Options{
 		Dialer: func(ctx context.Context) (net.Conn, error) {
 			var conn net.Conn
-			err := internal.WithSpan(ctx, "dialer", func(ctx context.Context) error {
+			err := internal.WithSpan(ctx, "dialer", func(ctx context.Context, span trace.Span) error {
 				var err error
-				trace.SpanFromContext(ctx).SetAttributes(
+				span.SetAttributes(
 					label.String("redis.network", opt.Network),
 					label.String("redis.addr", opt.Addr),
 				)
diff --git a/redis.go b/redis.go
index 3cff1b5..b70fef3 100644
--- a/redis.go
+++ b/redis.go
@@ -8,6 +8,8 @@ import (
 	"github.com/go-redis/redis/v8/internal"
 	"github.com/go-redis/redis/v8/internal/pool"
 	"github.com/go-redis/redis/v8/internal/proto"
+	"go.opentelemetry.io/otel/api/trace"
+	"go.opentelemetry.io/otel/label"
 )
 
 // Nil reply returned by Redis when key does not exist.
@@ -223,7 +225,7 @@ func (c *baseClient) _getConn(ctx context.Context) (*pool.Conn, error) {
 		return cn, nil
 	}
 
-	err = internal.WithSpan(ctx, "init_conn", func(ctx context.Context) error {
+	err = internal.WithSpan(ctx, "init_conn", func(ctx context.Context, span trace.Span) error {
 		return c.initConn(ctx, cn)
 	})
 	if err != nil {
@@ -297,11 +299,18 @@ 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, "with_conn", func(ctx context.Context) error {
+	return internal.WithSpan(ctx, "with_conn", func(ctx context.Context, span trace.Span) error {
 		cn, err := c.getConn(ctx)
 		if err != nil {
 			return err
 		}
+
+		if span.IsRecording() {
+			if remoteAddr := cn.RemoteAddr(); remoteAddr != nil {
+				span.SetAttributes(label.String("net.peer.ip", remoteAddr.String()))
+			}
+		}
+
 		defer func() {
 			c.releaseConn(ctx, cn, err)
 		}()
@@ -317,7 +326,7 @@ func (c *baseClient) process(ctx context.Context, cmd Cmder) error {
 		attempt := attempt
 
 		var retry bool
-		err := internal.WithSpan(ctx, "process", func(ctx context.Context) error {
+		err := internal.WithSpan(ctx, "process", func(ctx context.Context, span trace.Span) error {
 			if attempt > 0 {
 				if err := internal.Sleep(ctx, c.retryBackoff(attempt)); err != nil {
 					return err