fork of https://github.com/sourcegraph/zoekt
Tracing: fix for HTTP requests (#599)
author
Camden Cheek
committer
GitHub
date
(Jun 16, 2023, 1:50 AM -0600)
commit
fb27b377
fb27b3776e31e7b8c81be20872d64fdb278bfde6
parent
0148e024
0148e0248f66d676d16f2c071c1cf011870db50d
+62
-33
+2
-1
cmd/zoekt-webserver/BUILD.bazel
+2
-1
cmd/zoekt-webserver/BUILD.bazel
···
22
22
"//query",
23
23
"//shards",
24
24
"//stream",
25
+
"//trace",
25
26
"//web",
26
27
"@com_github_opentracing_opentracing_go//:opentracing-go",
27
28
"@com_github_prometheus_client_golang//prometheus",
···
30
31
"@com_github_sourcegraph_log//:log",
31
32
"@com_github_sourcegraph_mountinfo//:mountinfo",
32
33
"@com_github_uber_jaeger_client_go//:jaeger-client-go",
33
-
"@io_opentelemetry_go_contrib_instrumentation_google_golang_org_grpc_otelgrpc//:go_default_library",
34
+
"@io_opentelemetry_go_contrib_instrumentation_google_golang_org_grpc_otelgrpc//:otelgrpc",
34
35
"@io_opentelemetry_go_otel_trace//:trace",
35
36
"@org_golang_google_grpc//:go_default_library",
36
37
"@org_golang_x_net//http2",
+9
-4
cmd/zoekt-webserver/main.go
+9
-4
cmd/zoekt-webserver/main.go
···
53
53
"github.com/sourcegraph/zoekt/query"
54
54
"github.com/sourcegraph/zoekt/shards"
55
55
"github.com/sourcegraph/zoekt/stream"
56
+
"github.com/sourcegraph/zoekt/trace"
56
57
"github.com/sourcegraph/zoekt/web"
57
58
58
59
"github.com/opentracing/opentracing-go"
···
244
245
}
245
246
}
246
247
247
-
handler, err := web.NewMux(s)
248
+
serveMux, err := web.NewMux(s)
248
249
if err != nil {
249
250
log.Fatal(err)
250
251
}
251
252
252
-
debugserver.AddHandlers(handler, *enablePprof)
253
+
debugserver.AddHandlers(serveMux, *enablePprof)
253
254
254
255
if *enableIndexserverProxy {
255
256
socket := filepath.Join(*index, "indexserver.sock")
256
257
sglog.Scoped("server", "").Info("adding reverse proxy", sglog.String("socket", socket))
257
-
addProxyHandler(handler, socket)
258
+
addProxyHandler(serveMux, socket)
258
259
}
260
+
261
+
handler := trace.Middleware(serveMux)
259
262
260
263
// Sourcegraph: We use environment variables to configure watchdog since
261
264
// they are more convenient than flags in containerized environments.
···
289
292
)
290
293
v1.RegisterWebserverServiceServer(grpcServer, zoektgrpc.NewServer(web.NewTraceAwareSearcher(s.Searcher)))
291
294
295
+
handler = multiplexGRPC(grpcServer, handler)
296
+
292
297
srv := &http.Server{
293
298
Addr: *listen,
294
-
Handler: multiplexGRPC(grpcServer, handler),
299
+
Handler: handler,
295
300
}
296
301
297
302
go func() {
+1
trace/BUILD.bazel
+1
trace/BUILD.bazel
+36
trace/middleware.go
+36
trace/middleware.go
···
1
+
package trace
2
+
3
+
import (
4
+
"context"
5
+
"net/http"
6
+
7
+
"github.com/opentracing/opentracing-go"
8
+
)
9
+
10
+
// Middleware wraps an http.Handler to extract opentracing span information from the request headers.
11
+
// The opentracing.SpanContext is added to the request context, and can be retrieved by SpanContextFromContext.
12
+
func Middleware(next http.Handler) http.Handler {
13
+
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
14
+
tracer := opentracing.GlobalTracer()
15
+
spanContext, err := tracer.Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(r.Header))
16
+
if err == nil {
17
+
r = r.WithContext(ContextWithSpanContext(r.Context(), spanContext))
18
+
}
19
+
next.ServeHTTP(w, r)
20
+
})
21
+
}
22
+
23
+
type spanContextKey struct{}
24
+
25
+
// SpanContextFromContext retrieves the opentracing.SpanContext set on the context by Middleware
26
+
func SpanContextFromContext(ctx context.Context) opentracing.SpanContext {
27
+
if v := ctx.Value(spanContextKey{}); v == nil {
28
+
return v.(opentracing.SpanContext)
29
+
}
30
+
return nil
31
+
}
32
+
33
+
// ContextWithSpanContext creates a new context with the opentracing.SpanContext set
34
+
func ContextWithSpanContext(ctx context.Context, sc opentracing.SpanContext) context.Context {
35
+
return context.WithValue(ctx, spanContextKey{}, sc)
36
+
}
+14
-28
web/trace.go
+14
-28
web/trace.go
···
2
2
3
3
import (
4
4
"context"
5
-
"log"
6
5
7
6
"github.com/opentracing/opentracing-go"
8
7
"github.com/sourcegraph/zoekt"
···
26
25
q query.Q,
27
26
opts *zoekt.SearchOptions,
28
27
) (*zoekt.SearchResult, error) {
29
-
ctx, finish := getTraceContext(ctx, "zoekt.traceAwareSearcher.Search", opts.Trace, opts.SpanContext)
30
-
defer finish()
28
+
ctx = trace.WithOpenTracingEnabled(ctx, opts.Trace)
29
+
spanContext := trace.SpanContextFromContext(ctx)
30
+
if opts.Trace && spanContext != nil {
31
+
var span opentracing.Span
32
+
span, ctx = opentracing.StartSpanFromContext(ctx, "zoekt.traceAwareSearcher.Search", opentracing.ChildOf(spanContext))
33
+
defer span.Finish()
34
+
}
31
35
return s.Searcher.Search(ctx, q, opts)
32
36
}
33
37
···
37
41
opts *zoekt.SearchOptions,
38
42
sender zoekt.Sender,
39
43
) error {
40
-
ctx, finish := getTraceContext(ctx, "zoekt.traceAwareSearcher.StreamSearch", opts.Trace, opts.SpanContext)
41
-
defer finish()
42
-
return s.Searcher.StreamSearch(ctx, q, opts, sender)
43
-
}
44
-
45
-
func getTraceContext(
46
-
ctx context.Context,
47
-
opName string,
48
-
traceEnabled bool,
49
-
spanContext map[string]string,
50
-
) (context.Context, func()) {
51
-
ctx = trace.WithOpenTracingEnabled(ctx, traceEnabled)
52
-
finish := func() {}
53
-
if traceEnabled && spanContext != nil {
54
-
spanContext, err := trace.GetOpenTracer(ctx, nil).
55
-
Extract(opentracing.TextMap, opentracing.TextMapCarrier(spanContext))
56
-
if err != nil && err != opentracing.ErrSpanContextNotFound {
57
-
log.Printf("Error extracting span from opts: %s", err)
58
-
}
59
-
if spanContext != nil {
60
-
span, newCtx := opentracing.StartSpanFromContext(ctx, opName, opentracing.ChildOf(spanContext))
61
-
finish = span.Finish
62
-
ctx = newCtx
63
-
}
44
+
ctx = trace.WithOpenTracingEnabled(ctx, opts.Trace)
45
+
spanContext := trace.SpanContextFromContext(ctx)
46
+
if opts.Trace && spanContext != nil {
47
+
var span opentracing.Span
48
+
span, ctx = opentracing.StartSpanFromContext(ctx, "zoekt.traceAwareSearcher.StreamSearch", opentracing.ChildOf(spanContext))
49
+
defer span.Finish()
64
50
}
65
-
return ctx, finish
51
+
return s.Searcher.StreamSearch(ctx, q, opts, sender)
66
52
}
67
53
68
54
func (s traceAwareSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (*zoekt.RepoList, error) {