@@ -2,6 +2,7 @@ package rediscensus
2
2
3
3
import (
4
4
"context"
5
+ "net"
5
6
6
7
"go.opencensus.io/trace"
7
8
@@ -17,29 +18,54 @@ func NewTracingHook() *TracingHook {
17
18
return new (TracingHook )
18
19
}
19
20
20
- func (TracingHook ) BeforeProcess ( ctx context. Context , cmd redis.Cmder ) (context. Context , error ) {
21
- ctx , span := trace . StartSpan ( ctx , cmd . FullName ())
22
- span . AddAttributes ( trace .StringAttribute ( "db.system" , "redis" ),
23
- trace . StringAttribute ( "redis.cmd" , rediscmd . CmdString ( cmd )) )
21
+ func (TracingHook ) DialHook ( next redis.DialHook ) redis. DialHook {
22
+ return func ( ctx context. Context , network , addr string ) (net. Conn , error ) {
23
+ ctx , span := trace .StartSpan ( ctx , "dial" )
24
+ defer span . End ( )
24
25
25
- return ctx , nil
26
- }
26
+ span .AddAttributes (
27
+ trace .StringAttribute ("db.system" , "redis" ),
28
+ trace .StringAttribute ("network" , network ),
29
+ trace .StringAttribute ("addr" , addr ),
30
+ )
31
+
32
+ conn , err := next (ctx , network , addr )
33
+ if err != nil {
34
+ recordErrorOnOCSpan (ctx , span , err )
27
35
28
- func ( TracingHook ) AfterProcess ( ctx context. Context , cmd redis. Cmder ) error {
29
- span := trace . FromContext ( ctx )
30
- if err := cmd . Err (); err != nil {
31
- recordErrorOnOCSpan ( ctx , span , err )
36
+ return nil , err
37
+ }
38
+
39
+ return conn , nil
32
40
}
33
- span .End ()
34
- return nil
35
41
}
36
42
37
- func (TracingHook ) BeforeProcessPipeline (ctx context.Context , cmds []redis.Cmder ) (context.Context , error ) {
38
- return ctx , nil
43
+ func (TracingHook ) ProcessHook (next redis.ProcessHook ) redis.ProcessHook {
44
+ return func (ctx context.Context , cmd redis.Cmder ) error {
45
+ ctx , span := trace .StartSpan (ctx , cmd .FullName ())
46
+ defer span .End ()
47
+
48
+ span .AddAttributes (
49
+ trace .StringAttribute ("db.system" , "redis" ),
50
+ trace .StringAttribute ("redis.cmd" , rediscmd .CmdString (cmd )),
51
+ )
52
+
53
+ err := next (ctx , cmd )
54
+ if err != nil {
55
+ recordErrorOnOCSpan (ctx , span , err )
56
+ return err
57
+ }
58
+
59
+ if err = cmd .Err (); err != nil {
60
+ recordErrorOnOCSpan (ctx , span , err )
61
+ }
62
+
63
+ return nil
64
+ }
39
65
}
40
66
41
- func (TracingHook ) AfterProcessPipeline ( ctx context. Context , cmds [] redis.Cmder ) error {
42
- return nil
67
+ func (TracingHook ) ProcessPipelineHook ( next redis.ProcessPipelineHook ) redis. ProcessPipelineHook {
68
+ return next
43
69
}
44
70
45
71
func recordErrorOnOCSpan (ctx context.Context , span * trace.Span , err error ) {
0 commit comments