package zipkin
import (
"math/rand"
"net/http"
"strconv"
"golang.org/x/net/context"
"github.com/go-kit/kit/server"
)
// http://www.slideshare.net/johanoskarsson/zipkin-runtime-open-house
// https://groups.google.com/forum/#!topic/zipkin-user/KilwtSA0g1k
// https://gist.github.com/yoavaa/3478d3a0df666f21a98c
const (
// https://github.com/racker/tryfer#headers
traceIDHTTPHeader = "X-B3-TraceId"
spanIDHTTPHeader = "X-B3-SpanId"
parentSpanIDHTTPHeader = "X-B3-ParentSpanId"
clientSend = "cs"
serverReceive = "sr"
serverSend = "ss"
clientReceive = "cr"
)
// AnnotateEndpoint extracts a span from the context, adds server-receive and
// server-send annotations at the boundaries, and submits the span to the
// collector. If no span is present, a new span is generated and put in the
// context.
func AnnotateEndpoint(f func(int64, int64, int64) *Span, c Collector) func(server.Endpoint) server.Endpoint {
return func(e server.Endpoint) server.Endpoint {
return func(ctx context.Context, req server.Request) (server.Response, error) {
span, ctx := mustGetServerSpan(ctx, f)
span.Annotate(serverReceive)
defer func() { span.Annotate(serverSend); c.Collect(span) }()
return e(ctx, req)
}
}
}
// FromHTTP is a helper method that allows NewSpanFunc's factory function to
// be easily invoked by passing an HTTP request. The span name is the HTTP
// method. The trace, span, and parent span IDs are taken from the request
// headers.
func FromHTTP(f func(int64, int64, int64) *Span) func(*http.Request) *Span {
return func(r *http.Request) *Span {
return f(
getID(r.Header, traceIDHTTPHeader),
getID(r.Header, spanIDHTTPHeader),
getID(r.Header, parentSpanIDHTTPHeader),
)
}
}
// ToContext returns a function that satisfies transport/http.BeforeFunc. When
// invoked, it generates a Zipkin span from the incoming HTTP request, and
// saves it in the request context under the SpanContextKey.
func ToContext(f func(*http.Request) *Span) func(context.Context, *http.Request) context.Context {
return func(ctx context.Context, r *http.Request) context.Context {
return context.WithValue(ctx, SpanContextKey, f(r))
}
}
// NewChildSpan creates a new child (client) span. If a span is present in the
// context, it will be interpreted as the parent.
func NewChildSpan(ctx context.Context, f func(int64, int64, int64) *Span) *Span {
val := ctx.Value(SpanContextKey)
if val == nil {
return f(newID(), newID(), 0)
}
parentSpan, ok := val.(*Span)
if !ok {
panic(SpanContextKey + " value isn't a span object")
}
var (
traceID = parentSpan.TraceID()
spanID = newID()
parentSpanID = parentSpan.SpanID()
)
return f(traceID, spanID, parentSpanID)
}
// SetRequestHeaders sets up HTTP headers for a new outbound request based on
// the (client) span. All IDs are encoded as hex strings.
func SetRequestHeaders(h http.Header, s *Span) {
if id := s.TraceID(); id > 0 {
h.Set(traceIDHTTPHeader, strconv.FormatInt(id, 16))
}
if id := s.SpanID(); id > 0 {
h.Set(spanIDHTTPHeader, strconv.FormatInt(id, 16))
}
if id := s.ParentSpanID(); id > 0 {
h.Set(parentSpanIDHTTPHeader, strconv.FormatInt(id, 16))
}
}
func mustGetServerSpan(ctx context.Context, f func(int64, int64, int64) *Span) (*Span, context.Context) {
val := ctx.Value(SpanContextKey)
if val == nil {
span := f(newID(), newID(), 0)
return span, context.WithValue(ctx, SpanContextKey, span)
}
span, ok := val.(*Span)
if !ok {
panic(SpanContextKey + " value isn't a span object")
}
return span, ctx
}
func getID(h http.Header, key string) int64 {
val := h.Get(key)
if val == "" {
return 0
}
i, err := strconv.ParseInt(val, 16, 64)
if err != nil {
panic("invalid Zipkin ID in HTTP header: " + val)
}
return i
}
func newID() int64 {
// https://github.com/wadey/go-zipkin/blob/46e5f01/trace.go#L183-188
// https://github.com/twitter/zipkin/issues/199
// :(
return rand.Int63() & 0x001fffffffffffff
}