1// Copyright 2019, OpenTelemetry Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package testtrace
16
17import (
18	"context"
19	"sync"
20	"time"
21
22	"go.opentelemetry.io/otel/api/core"
23	"go.opentelemetry.io/otel/api/trace"
24)
25
26var _ trace.Tracer = (*Tracer)(nil)
27
28// Tracer is a type of OpenTelemetry Tracer that tracks both active and ended spans,
29// and which creates Spans that may be inspected to see what data has been set on them.
30type Tracer struct {
31	lock      *sync.RWMutex
32	generator Generator
33	spans     []*Span
34}
35
36func NewTracer(opts ...TracerOption) *Tracer {
37	c := newTracerConfig(opts...)
38
39	return &Tracer{
40		lock:      &sync.RWMutex{},
41		generator: c.generator,
42	}
43}
44
45func (t *Tracer) Start(ctx context.Context, name string, opts ...trace.StartOption) (context.Context, trace.Span) {
46	var c trace.StartConfig
47
48	for _, opt := range opts {
49		opt(&c)
50	}
51
52	var traceID core.TraceID
53	var parentSpanID core.SpanID
54
55	if parentSpanContext := c.Relation.SpanContext; parentSpanContext.IsValid() {
56		traceID = parentSpanContext.TraceID
57		parentSpanID = parentSpanContext.SpanID
58	} else if parentSpanContext := trace.SpanFromContext(ctx).SpanContext(); parentSpanContext.IsValid() {
59		traceID = parentSpanContext.TraceID
60		parentSpanID = parentSpanContext.SpanID
61	} else {
62		traceID = t.generator.TraceID()
63	}
64
65	spanID := t.generator.SpanID()
66
67	startTime := time.Now()
68
69	if st := c.StartTime; !st.IsZero() {
70		startTime = st
71	}
72
73	span := &Span{
74		lock:      &sync.RWMutex{},
75		tracer:    t,
76		startTime: startTime,
77		spanContext: core.SpanContext{
78			TraceID: traceID,
79			SpanID:  spanID,
80		},
81		parentSpanID: parentSpanID,
82		attributes:   make(map[core.Key]core.Value),
83		links:        make(map[core.SpanContext][]core.KeyValue),
84	}
85
86	span.SetName(name)
87	span.SetAttributes(c.Attributes...)
88
89	for _, link := range c.Links {
90		span.links[link.SpanContext] = link.Attributes
91	}
92
93	t.lock.Lock()
94
95	t.spans = append(t.spans, span)
96
97	t.lock.Unlock()
98
99	return trace.ContextWithSpan(ctx, span), span
100}
101
102func (t *Tracer) WithSpan(ctx context.Context, name string, body func(ctx context.Context) error) error {
103	ctx, _ = t.Start(ctx, name)
104
105	return body(ctx)
106}
107
108// Spans returns the list of current and ended Spans started via the Tracer.
109func (t *Tracer) Spans() []*Span {
110	t.lock.RLock()
111	defer t.lock.RUnlock()
112
113	return append([]*Span{}, t.spans...)
114}
115
116// TracerOption enables configuration of a new Tracer.
117type TracerOption func(*tracerConfig)
118
119// TracerWithGenerator enables customization of the Generator that the Tracer will use
120// to create new trace and span IDs.
121// By default, new Tracers will use the CountGenerator.
122func TracerWithGenerator(generator Generator) TracerOption {
123	return func(c *tracerConfig) {
124		c.generator = generator
125	}
126}
127
128type tracerConfig struct {
129	generator Generator
130}
131
132func newTracerConfig(opts ...TracerOption) tracerConfig {
133	var c tracerConfig
134	defaultOpts := []TracerOption{
135		TracerWithGenerator(NewCountGenerator()),
136	}
137
138	for _, opt := range append(defaultOpts, opts...) {
139		opt(&c)
140	}
141
142	return c
143}
144