fork of https://github.com/sourcegraph/zoekt
1package web
2
3import (
4 "context"
5 "log"
6
7 "github.com/opentracing/opentracing-go"
8 "github.com/sourcegraph/zoekt"
9 "github.com/sourcegraph/zoekt/query"
10 "github.com/sourcegraph/zoekt/trace"
11)
12
13// traceAwareSearcher wraps a zoekt.Searcher instance so that the tracing context item is set in the
14// context. This context item toggles on trace collection via the
15// github.com/sourcegraph/zoekt/trace/ot package.
16type traceAwareSearcher struct {
17 Searcher zoekt.Streamer
18}
19
20func (s traceAwareSearcher) Search(
21 ctx context.Context,
22 q query.Q,
23 opts *zoekt.SearchOptions,
24) (*zoekt.SearchResult, error) {
25 ctx, finish := getTraceContext(ctx, "zoekt.traceAwareSearcher.Search", opts.Trace, opts.SpanContext)
26 defer finish()
27 return s.Searcher.Search(ctx, q, opts)
28}
29
30func (s traceAwareSearcher) StreamSearch(
31 ctx context.Context,
32 q query.Q,
33 opts *zoekt.SearchOptions,
34 sender zoekt.Sender,
35) error {
36 ctx, finish := getTraceContext(ctx, "zoekt.traceAwareSearcher.StreamSearch", opts.Trace, opts.SpanContext)
37 defer finish()
38 return s.Searcher.StreamSearch(ctx, q, opts, sender)
39}
40
41func getTraceContext(
42 ctx context.Context,
43 opName string,
44 traceEnabled bool,
45 spanContext map[string]string,
46) (context.Context, func()) {
47 ctx = trace.WithOpenTracingEnabled(ctx, traceEnabled)
48 finish := func() {}
49 if traceEnabled && spanContext != nil {
50 spanContext, err := trace.GetOpenTracer(ctx, nil).
51 Extract(opentracing.TextMap, opentracing.TextMapCarrier(spanContext))
52 if err != nil && err != opentracing.ErrSpanContextNotFound {
53 log.Printf("Error extracting span from opts: %s", err)
54 }
55 if spanContext != nil {
56 span, newCtx := opentracing.StartSpanFromContext(ctx, opName, opentracing.ChildOf(spanContext))
57 finish = span.Finish
58 ctx = newCtx
59 }
60 }
61 return ctx, finish
62}
63
64func (s traceAwareSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (*zoekt.RepoList, error) {
65 return s.Searcher.List(ctx, q, opts)
66}
67func (s traceAwareSearcher) Close() { s.Searcher.Close() }
68func (s traceAwareSearcher) String() string { return s.Searcher.String() }