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