67 lines
1.6 KiB
Go
67 lines
1.6 KiB
Go
package tracing
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"github.com/urfave/negroni"
|
|
)
|
|
|
|
// NegroniHandlerWrapper is used to wrap negroni handler middleware
|
|
type NegroniHandlerWrapper struct {
|
|
name string
|
|
next negroni.Handler
|
|
clientSpanKind bool
|
|
}
|
|
|
|
// HTTPHandlerWrapper is used to wrap http handler middleware
|
|
type HTTPHandlerWrapper struct {
|
|
name string
|
|
handler http.Handler
|
|
clientSpanKind bool
|
|
}
|
|
|
|
// NewNegroniHandlerWrapper return a negroni.Handler struct
|
|
func (t *Tracing) NewNegroniHandlerWrapper(name string, handler negroni.Handler, clientSpanKind bool) negroni.Handler {
|
|
if t.IsEnabled() && handler != nil {
|
|
return &NegroniHandlerWrapper{
|
|
name: name,
|
|
next: handler,
|
|
clientSpanKind: clientSpanKind,
|
|
}
|
|
}
|
|
return handler
|
|
}
|
|
|
|
// NewHTTPHandlerWrapper return a http.Handler struct
|
|
func (t *Tracing) NewHTTPHandlerWrapper(name string, handler http.Handler, clientSpanKind bool) http.Handler {
|
|
if t.IsEnabled() && handler != nil {
|
|
return &HTTPHandlerWrapper{
|
|
name: name,
|
|
handler: handler,
|
|
clientSpanKind: clientSpanKind,
|
|
}
|
|
}
|
|
return handler
|
|
}
|
|
|
|
func (t *NegroniHandlerWrapper) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
|
|
var finish func()
|
|
_, r, finish = StartSpan(r, t.name, t.clientSpanKind)
|
|
defer finish()
|
|
|
|
if t.next != nil {
|
|
t.next.ServeHTTP(rw, r, next)
|
|
}
|
|
}
|
|
|
|
func (t *HTTPHandlerWrapper) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
|
|
var finish func()
|
|
_, r, finish = StartSpan(r, t.name, t.clientSpanKind)
|
|
defer finish()
|
|
|
|
if t.handler != nil {
|
|
t.handler.ServeHTTP(rw, r)
|
|
}
|
|
|
|
}
|