1// Copyright The 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 internal
16
17import (
18	"context"
19	"math/rand"
20	"reflect"
21	"sync"
22	"time"
23
24	"go.opentelemetry.io/otel/codes"
25	"go.opentelemetry.io/otel/internal/baggage"
26	otelparent "go.opentelemetry.io/otel/internal/trace/parent"
27	"go.opentelemetry.io/otel/label"
28	"go.opentelemetry.io/otel/trace"
29
30	"go.opentelemetry.io/otel/bridge/opentracing/migration"
31)
32
33var (
34	ComponentKey     = label.Key("component")
35	ServiceKey       = label.Key("service")
36	StatusCodeKey    = label.Key("status.code")
37	StatusMessageKey = label.Key("status.message")
38	ErrorKey         = label.Key("error")
39	NameKey          = label.Key("name")
40)
41
42type MockContextKeyValue struct {
43	Key   interface{}
44	Value interface{}
45}
46
47type MockTracer struct {
48	Resources             baggage.Map
49	FinishedSpans         []*MockSpan
50	SpareTraceIDs         []trace.TraceID
51	SpareSpanIDs          []trace.SpanID
52	SpareContextKeyValues []MockContextKeyValue
53
54	randLock sync.Mutex
55	rand     *rand.Rand
56}
57
58var _ trace.Tracer = &MockTracer{}
59var _ migration.DeferredContextSetupTracerExtension = &MockTracer{}
60
61func NewMockTracer() *MockTracer {
62	return &MockTracer{
63		Resources:             baggage.NewEmptyMap(),
64		FinishedSpans:         nil,
65		SpareTraceIDs:         nil,
66		SpareSpanIDs:          nil,
67		SpareContextKeyValues: nil,
68
69		rand: rand.New(rand.NewSource(time.Now().Unix())),
70	}
71}
72
73func (t *MockTracer) Start(ctx context.Context, name string, opts ...trace.SpanOption) (context.Context, trace.Span) {
74	config := trace.NewSpanConfig(opts...)
75	startTime := config.Timestamp
76	if startTime.IsZero() {
77		startTime = time.Now()
78	}
79	spanContext := trace.SpanContext{
80		TraceID:    t.getTraceID(ctx, config),
81		SpanID:     t.getSpanID(),
82		TraceFlags: 0,
83	}
84	span := &MockSpan{
85		mockTracer:     t,
86		officialTracer: t,
87		spanContext:    spanContext,
88		recording:      config.Record,
89		Attributes: baggage.NewMap(baggage.MapUpdate{
90			MultiKV: config.Attributes,
91		}),
92		StartTime:    startTime,
93		EndTime:      time.Time{},
94		ParentSpanID: t.getParentSpanID(ctx, config),
95		Events:       nil,
96		SpanKind:     trace.ValidateSpanKind(config.SpanKind),
97	}
98	if !migration.SkipContextSetup(ctx) {
99		ctx = trace.ContextWithSpan(ctx, span)
100		ctx = t.addSpareContextValue(ctx)
101	}
102	return ctx, span
103}
104
105func (t *MockTracer) addSpareContextValue(ctx context.Context) context.Context {
106	if len(t.SpareContextKeyValues) > 0 {
107		pair := t.SpareContextKeyValues[0]
108		t.SpareContextKeyValues[0] = MockContextKeyValue{}
109		t.SpareContextKeyValues = t.SpareContextKeyValues[1:]
110		if len(t.SpareContextKeyValues) == 0 {
111			t.SpareContextKeyValues = nil
112		}
113		ctx = context.WithValue(ctx, pair.Key, pair.Value)
114	}
115	return ctx
116}
117
118func (t *MockTracer) getTraceID(ctx context.Context, config *trace.SpanConfig) trace.TraceID {
119	if parent := t.getParentSpanContext(ctx, config); parent.IsValid() {
120		return parent.TraceID
121	}
122	if len(t.SpareTraceIDs) > 0 {
123		traceID := t.SpareTraceIDs[0]
124		t.SpareTraceIDs = t.SpareTraceIDs[1:]
125		if len(t.SpareTraceIDs) == 0 {
126			t.SpareTraceIDs = nil
127		}
128		return traceID
129	}
130	return t.getRandTraceID()
131}
132
133func (t *MockTracer) getParentSpanID(ctx context.Context, config *trace.SpanConfig) trace.SpanID {
134	if parent := t.getParentSpanContext(ctx, config); parent.IsValid() {
135		return parent.SpanID
136	}
137	return trace.SpanID{}
138}
139
140func (t *MockTracer) getParentSpanContext(ctx context.Context, config *trace.SpanConfig) trace.SpanContext {
141	spanCtx, _, _ := otelparent.GetSpanContextAndLinks(ctx, config.NewRoot)
142	return spanCtx
143}
144
145func (t *MockTracer) getSpanID() trace.SpanID {
146	if len(t.SpareSpanIDs) > 0 {
147		spanID := t.SpareSpanIDs[0]
148		t.SpareSpanIDs = t.SpareSpanIDs[1:]
149		if len(t.SpareSpanIDs) == 0 {
150			t.SpareSpanIDs = nil
151		}
152		return spanID
153	}
154	return t.getRandSpanID()
155}
156
157func (t *MockTracer) getRandSpanID() trace.SpanID {
158	t.randLock.Lock()
159	defer t.randLock.Unlock()
160
161	sid := trace.SpanID{}
162	t.rand.Read(sid[:])
163
164	return sid
165}
166
167func (t *MockTracer) getRandTraceID() trace.TraceID {
168	t.randLock.Lock()
169	defer t.randLock.Unlock()
170
171	tid := trace.TraceID{}
172	t.rand.Read(tid[:])
173
174	return tid
175}
176
177func (t *MockTracer) DeferredContextSetupHook(ctx context.Context, span trace.Span) context.Context {
178	return t.addSpareContextValue(ctx)
179}
180
181type MockEvent struct {
182	Timestamp  time.Time
183	Name       string
184	Attributes baggage.Map
185}
186
187type MockSpan struct {
188	mockTracer     *MockTracer
189	officialTracer trace.Tracer
190	spanContext    trace.SpanContext
191	SpanKind       trace.SpanKind
192	recording      bool
193
194	Attributes   baggage.Map
195	StartTime    time.Time
196	EndTime      time.Time
197	ParentSpanID trace.SpanID
198	Events       []MockEvent
199}
200
201var _ trace.Span = &MockSpan{}
202var _ migration.OverrideTracerSpanExtension = &MockSpan{}
203
204func (s *MockSpan) SpanContext() trace.SpanContext {
205	return s.spanContext
206}
207
208func (s *MockSpan) IsRecording() bool {
209	return s.recording
210}
211
212func (s *MockSpan) SetStatus(code codes.Code, msg string) {
213	s.SetAttributes(StatusCodeKey.Uint32(uint32(code)), StatusMessageKey.String(msg))
214}
215
216func (s *MockSpan) SetName(name string) {
217	s.SetAttributes(NameKey.String(name))
218}
219
220func (s *MockSpan) SetError(v bool) {
221	s.SetAttributes(ErrorKey.Bool(v))
222}
223
224func (s *MockSpan) SetAttributes(attributes ...label.KeyValue) {
225	s.applyUpdate(baggage.MapUpdate{
226		MultiKV: attributes,
227	})
228}
229
230func (s *MockSpan) applyUpdate(update baggage.MapUpdate) {
231	s.Attributes = s.Attributes.Apply(update)
232}
233
234func (s *MockSpan) End(options ...trace.SpanOption) {
235	if !s.EndTime.IsZero() {
236		return // already finished
237	}
238	config := trace.NewSpanConfig(options...)
239	endTime := config.Timestamp
240	if endTime.IsZero() {
241		endTime = time.Now()
242	}
243	s.EndTime = endTime
244	s.mockTracer.FinishedSpans = append(s.mockTracer.FinishedSpans, s)
245}
246
247func (s *MockSpan) RecordError(err error, opts ...trace.EventOption) {
248	if err == nil {
249		return // no-op on nil error
250	}
251
252	if !s.EndTime.IsZero() {
253		return // already finished
254	}
255
256	s.SetStatus(codes.Error, "")
257	opts = append(opts, trace.WithAttributes(
258		label.String("error.type", reflect.TypeOf(err).String()),
259		label.String("error.message", err.Error()),
260	))
261	s.AddEvent("error", opts...)
262}
263
264func (s *MockSpan) Tracer() trace.Tracer {
265	return s.officialTracer
266}
267
268func (s *MockSpan) AddEvent(name string, o ...trace.EventOption) {
269	c := trace.NewEventConfig(o...)
270	s.Events = append(s.Events, MockEvent{
271		Timestamp: c.Timestamp,
272		Name:      name,
273		Attributes: baggage.NewMap(baggage.MapUpdate{
274			MultiKV: c.Attributes,
275		}),
276	})
277}
278
279func (s *MockSpan) OverrideTracer(tracer trace.Tracer) {
280	s.officialTracer = tracer
281}
282