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