From 805bfc2c60546190e6c1a160a279dd6ebe2d6de4 Mon Sep 17 00:00:00 2001
From: Daniel Cormier <dcormier@users.noreply.github.com>
Date: Thu, 26 Jan 2023 01:23:14 -0500
Subject: [PATCH] fix: implement `redis.Hook` (v9) for
 `rediscensus.TracingHook` (#2367)

Fixes #2266
---
 extra/rediscensus/rediscensus.go | 58 +++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/extra/rediscensus/rediscensus.go b/extra/rediscensus/rediscensus.go
index 4fdb0a99..bc48608a 100644
--- a/extra/rediscensus/rediscensus.go
+++ b/extra/rediscensus/rediscensus.go
@@ -2,6 +2,7 @@ package rediscensus
 
 import (
 	"context"
+	"net"
 
 	"go.opencensus.io/trace"
 
@@ -17,29 +18,54 @@ func NewTracingHook() *TracingHook {
 	return new(TracingHook)
 }
 
-func (TracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
-	ctx, span := trace.StartSpan(ctx, cmd.FullName())
-	span.AddAttributes(trace.StringAttribute("db.system", "redis"),
-		trace.StringAttribute("redis.cmd", rediscmd.CmdString(cmd)))
+func (TracingHook) DialHook(next redis.DialHook) redis.DialHook {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		ctx, span := trace.StartSpan(ctx, "dial")
+		defer span.End()
 
-	return ctx, nil
-}
+		span.AddAttributes(
+			trace.StringAttribute("db.system", "redis"),
+			trace.StringAttribute("network", network),
+			trace.StringAttribute("addr", addr),
+		)
 
-func (TracingHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error {
-	span := trace.FromContext(ctx)
-	if err := cmd.Err(); err != nil {
-		recordErrorOnOCSpan(ctx, span, err)
+		conn, err := next(ctx, network, addr)
+		if err != nil {
+			recordErrorOnOCSpan(ctx, span, err)
+
+			return nil, err
+		}
+
+		return conn, nil
 	}
-	span.End()
-	return nil
 }
 
-func (TracingHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) {
-	return ctx, nil
+func (TracingHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
+	return func(ctx context.Context, cmd redis.Cmder) error {
+		ctx, span := trace.StartSpan(ctx, cmd.FullName())
+		defer span.End()
+
+		span.AddAttributes(
+			trace.StringAttribute("db.system", "redis"),
+			trace.StringAttribute("redis.cmd", rediscmd.CmdString(cmd)),
+		)
+
+		err := next(ctx, cmd)
+		if err != nil {
+			recordErrorOnOCSpan(ctx, span, err)
+			return err
+		}
+
+		if err = cmd.Err(); err != nil {
+			recordErrorOnOCSpan(ctx, span, err)
+		}
+
+		return nil
+	}
 }
 
-func (TracingHook) AfterProcessPipeline(ctx context.Context, cmds []redis.Cmder) error {
-	return nil
+func (TracingHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
+	return next
 }
 
 func recordErrorOnOCSpan(ctx context.Context, span *trace.Span, err error) {