forked from coredns/coredns
-
Notifications
You must be signed in to change notification settings - Fork 0
/
trace.go
90 lines (77 loc) · 1.93 KB
/
trace.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
// Package trace implements OpenTracing-based tracing
package trace
import (
"fmt"
"sync"
"sync/atomic"
"github.com/coredns/coredns/middleware"
"github.com/miekg/dns"
ot "github.com/opentracing/opentracing-go"
zipkin "github.com/openzipkin/zipkin-go-opentracing"
"golang.org/x/net/context"
)
// Trace holds the tracer and endpoint info
type Trace interface {
middleware.Handler
Tracer() ot.Tracer
}
type trace struct {
Next middleware.Handler
ServiceEndpoint string
Endpoint string
EndpointType string
tracer ot.Tracer
serviceName string
clientServer bool
every uint64
count uint64
Once sync.Once
}
func (t *trace) Tracer() ot.Tracer {
return t.tracer
}
// OnStartup sets up the tracer
func (t *trace) OnStartup() error {
var err error
t.Once.Do(func() {
switch t.EndpointType {
case "zipkin":
err = t.setupZipkin()
default:
err = fmt.Errorf("Unknown endpoint type: %s", t.EndpointType)
}
})
return err
}
func (t *trace) setupZipkin() error {
collector, err := zipkin.NewHTTPCollector(t.Endpoint)
if err != nil {
return err
}
recorder := zipkin.NewRecorder(collector, false, t.ServiceEndpoint, t.serviceName)
t.tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(t.clientServer))
if err != nil {
return err
}
return nil
}
// Name implements the Handler interface.
func (t *trace) Name() string {
return "trace"
}
// ServeDNS implements the middleware.Handle interface.
func (t *trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
trace := false
if t.every > 0 {
queryNr := atomic.AddUint64(&t.count, 1)
if queryNr%t.every == 0 {
trace = true
}
}
if span := ot.SpanFromContext(ctx); span == nil && trace {
span := t.Tracer().StartSpan("servedns")
defer span.Finish()
ctx = ot.ContextWithSpan(ctx, span)
}
return middleware.NextOrFailure(t.Name(), t.Next, ctx, w, r)
}