Skip to content

Commit caf28f5

Browse files
tonyredondotonyredondo
authored
Merge branch 'main' into anmarchenko/tests_skipped_tag_propagation
2 parents e91803f + 5fd6d9c commit caf28f5

13 files changed

Lines changed: 980 additions & 87 deletions

‎ddtrace/mocktracer/civisibilitymocktracer.go‎

Lines changed: 157 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -17,31 +17,88 @@ import (
1717
)
1818

1919
type civisibilitymocktracer struct {
20-
mock *mocktracer // mock tracer
21-
real tracer.Tracer // real tracer (for the testotimization/civisibility spans)
20+
// mock records user-created spans so tests that use mocktracer keep their existing assertions.
21+
mock *mocktracer
22+
23+
// realMu protects real while CI Visibility startup, span creation, flush, and shutdown can overlap.
24+
realMu sync.RWMutex
25+
real tracer.Tracer // real receives CI Visibility spans that must not be captured by mock.
26+
27+
// realSpansMu protects realSpans, which tracks spans that must bypass mock FinishSpan handling.
28+
realSpansMu sync.Mutex
29+
realSpans map[*tracer.Span]struct{}
30+
31+
// isnoop disables user-facing mock behavior after Stop while allowing CI Visibility cleanup to continue.
2232
isnoop atomic.Bool
2333
}
2434

2535
var (
2636
_ tracer.Tracer = (*civisibilitymocktracer)(nil)
2737
_ Tracer = (*civisibilitymocktracer)(nil)
28-
29-
realSpans = make(map[*tracer.Span]bool)
30-
realSpansMutex sync.Mutex
3138
)
3239

33-
// Creates a new CIVisibilityMockTracer that uses the mock tracer for all spans except the CIVisibility spans.
40+
// newCIVisibilityMockTracer creates a mock tracer that delegates CI Visibility spans to the real tracer.
3441
func newCIVisibilityMockTracer() *civisibilitymocktracer {
3542
currentTracer := getGlobalTracer()
36-
// let's check if the current tracer is already a civisibilitymocktracer
37-
// if so, we need to get the real tracer from it
43+
// Repeated mocktracer starts should unwrap the previous CI Visibility mock tracer
44+
// and keep its real tracer delegate instead of stacking wrappers.
3845
if currentCIVisibilityMockTracer, ok := currentTracer.(*civisibilitymocktracer); ok && currentCIVisibilityMockTracer != nil {
39-
currentTracer = currentCIVisibilityMockTracer.real
46+
currentTracer = currentCIVisibilityMockTracer.realTracer()
4047
}
4148
return &civisibilitymocktracer{
42-
mock: newMockTracer(),
43-
real: currentTracer,
49+
mock: newMockTracer(),
50+
real: currentTracer,
51+
realSpans: make(map[*tracer.Span]struct{}),
52+
}
53+
}
54+
55+
// realTracer returns the currently installed CI Visibility tracer delegate.
56+
func (t *civisibilitymocktracer) realTracer() tracer.Tracer {
57+
t.realMu.RLock()
58+
defer t.realMu.RUnlock()
59+
return t.real
60+
}
61+
62+
// SetCIVisibilityTracer installs the tracer used for CI Visibility spans while
63+
// keeping this mock tracer as the process global tracer.
64+
func (t *civisibilitymocktracer) SetCIVisibilityTracer(real tracer.Tracer) bool {
65+
if real == nil {
66+
return false
67+
}
68+
69+
t.realMu.Lock()
70+
if t.isnoop.Load() {
71+
old := t.real
72+
t.real = &tracer.NoopTracer{}
73+
t.realMu.Unlock()
74+
if old != nil && old != real {
75+
stopRealTracerDelegate(old)
76+
}
77+
return false
78+
}
79+
old := t.real
80+
t.real = real
81+
t.realMu.Unlock()
82+
83+
if old != nil && old != real {
84+
stopRealTracerDelegate(old)
85+
}
86+
return true
87+
}
88+
89+
// stopRealTracerDelegate stops a tracer owned by civisibilitymocktracer without
90+
// letting mocktracer cleanup overwrite the process global tracer.
91+
func stopRealTracerDelegate(real tracer.Tracer) {
92+
if real == nil {
93+
return
94+
}
95+
if mt, ok := real.(*mocktracer); ok {
96+
if mt.dsmProcessor != nil {
97+
mt.dsmProcessor.Stop()
98+
}
99+
return
44100
}
101+
real.Stop()
45102
}
46103

47104
// SentDSMBacklogs returns the Data Streams Monitoring backlogs that have been sent by the mock tracer.
@@ -55,38 +112,53 @@ func (t *civisibilitymocktracer) SentDSMBacklogs() []datastreams.Backlog {
55112
return t.mock.dsmTransport.backlogs
56113
}
57114

58-
// Stop deactivates the CIVisibility mock tracer by setting it to noop mode and stopping
59-
// the Data Streams Monitoring processor. This should be called when testing has finished.
115+
// Stop deactivates the CI Visibility mock tracer by setting it to noop mode and stopping
116+
// the Data Streams Monitoring processor. If this wrapper has already been removed from
117+
// the global tracer slot, it also stops the real delegate because CI Visibility shutdown
118+
// can no longer reach it through the global tracer.
60119
func (t *civisibilitymocktracer) Stop() {
120+
var realToStop tracer.Tracer
121+
removedFromGlobalTracer := getGlobalTracer() != t
122+
123+
t.realMu.Lock()
61124
t.isnoop.Store(true)
62-
t.mock.dsmProcessor.Stop()
63-
if civisibility.GetState() == civisibility.StateExiting {
64-
t.real.Stop()
125+
if removedFromGlobalTracer || civisibility.GetState() == civisibility.StateExiting {
126+
realToStop = t.real
65127
t.real = &tracer.NoopTracer{}
66128
}
129+
t.realMu.Unlock()
130+
131+
t.mock.dsmProcessor.Stop()
132+
stopRealTracerDelegate(realToStop)
67133
}
68134

69135
// StartSpan creates a new span with the given operation name and options. If the span type
70136
// indicates it's a CI Visibility span (like a test session, module, suite, or individual test),
71137
// it uses the real tracer to create the span. For all other spans, it uses the mock tracer.
72-
// If the tracer is in noop mode, it returns nil.
138+
// If the mock tracer is in noop mode, non-CI Visibility spans return nil while
139+
// CI Visibility spans may still use the real tracer until CI Visibility exits.
73140
func (t *civisibilitymocktracer) StartSpan(operationName string, opts ...tracer.StartSpanOption) *tracer.Span {
74-
if t.real != nil {
75-
var cfg tracer.StartSpanConfig
76-
for _, fn := range opts {
77-
fn(&cfg)
78-
}
141+
var cfg tracer.StartSpanConfig
142+
for _, fn := range opts {
143+
fn(&cfg)
144+
}
145+
146+
if isCIVisibilitySpan(cfg) {
147+
t.realMu.RLock()
148+
real := t.real
149+
if real != nil {
150+
// If the span is a CI Visibility span, use the real tracer to create it.
151+
realSpan := real.StartSpan(operationName, opts...)
152+
t.realMu.RUnlock()
79153

80-
if spanType, ok := cfg.Tags[ext.SpanType]; ok &&
81-
(spanType == constants.SpanTypeTestSession || spanType == constants.SpanTypeTestModule ||
82-
spanType == constants.SpanTypeTestSuite || spanType == constants.SpanTypeTest) {
83-
// If the span is a civisibility span, use the real tracer to create it.
84-
realSpan := t.real.StartSpan(operationName, opts...)
85-
realSpansMutex.Lock()
86-
defer realSpansMutex.Unlock()
87-
realSpans[realSpan] = true
154+
if realSpan != nil {
155+
t.realSpansMu.Lock()
156+
t.realSpans[realSpan] = struct{}{}
157+
t.realSpansMu.Unlock()
158+
}
88159
return realSpan
89160
}
161+
t.realMu.RUnlock()
90162
}
91163

92164
if t.isnoop.Load() {
@@ -97,14 +169,28 @@ func (t *civisibilitymocktracer) StartSpan(operationName string, opts ...tracer.
97169
return t.mock.StartSpan(operationName, opts...)
98170
}
99171

100-
// FinishSpan marks the given span as finished in the mock tracer. This is called by spans
101-
// when they finish, adding them to the list of finished spans for later inspection.
172+
// isCIVisibilitySpan reports whether cfg describes a CI Visibility span that
173+
// must bypass user-facing mocktracer storage.
174+
func isCIVisibilitySpan(cfg tracer.StartSpanConfig) bool {
175+
spanType, ok := cfg.Tags[ext.SpanType]
176+
return ok && (spanType == constants.SpanTypeTestSession ||
177+
spanType == constants.SpanTypeTestModule ||
178+
spanType == constants.SpanTypeTestSuite ||
179+
spanType == constants.SpanTypeTest)
180+
}
181+
182+
// FinishSpan marks mock-created spans as finished while keeping CI Visibility
183+
// spans out of the user-facing mock span list.
102184
func (t *civisibilitymocktracer) FinishSpan(s *tracer.Span) {
103-
realSpansMutex.Lock()
104-
defer realSpansMutex.Unlock()
185+
if s == nil {
186+
return
187+
}
188+
189+
t.realSpansMu.Lock()
105190
// Check if the span is a real span (i.e., created by the real tracer).
106-
if _, isRealSpan := realSpans[s]; isRealSpan {
107-
delete(realSpans, s)
191+
_, isRealSpan := t.realSpans[s]
192+
t.realSpansMu.Unlock()
193+
if isRealSpan {
108194
return
109195
}
110196
if t.isnoop.Load() {
@@ -113,6 +199,31 @@ func (t *civisibilitymocktracer) FinishSpan(s *tracer.Span) {
113199
t.mock.FinishSpan(s)
114200
}
115201

202+
// TracerForFinishedChunk returns the current real tracer when a finished chunk
203+
// contains CI Visibility spans created by this wrapper.
204+
func (t *civisibilitymocktracer) TracerForFinishedChunk(spans []*tracer.Span) (tracer.Tracer, bool) {
205+
hasRealSpan := false
206+
t.realSpansMu.Lock()
207+
for _, s := range spans {
208+
if _, ok := t.realSpans[s]; ok {
209+
delete(t.realSpans, s)
210+
hasRealSpan = true
211+
}
212+
}
213+
t.realSpansMu.Unlock()
214+
if !hasRealSpan {
215+
return nil, false
216+
}
217+
218+
t.realMu.RLock()
219+
real := t.real
220+
t.realMu.RUnlock()
221+
if real == nil {
222+
return nil, false
223+
}
224+
return real, true
225+
}
226+
116227
// GetDataStreamsProcessor returns the Data Streams Monitoring processor used by the mock tracer.
117228
// If the tracer is in noop mode, it returns nil. This processor is used to monitor
118229
// and record data stream metrics.
@@ -163,12 +274,21 @@ func (t *civisibilitymocktracer) Inject(context *tracer.SpanContext, carrier any
163274
}
164275

165276
func (t *civisibilitymocktracer) TracerConf() tracer.TracerConf {
277+
t.realMu.RLock()
278+
defer t.realMu.RUnlock()
279+
if t.real == nil {
280+
return tracer.TracerConf{}
281+
}
166282
return t.real.TracerConf()
167283
}
168284

169285
// Flush forces a flush of both the mock tracer and the real tracer.
170286
// This ensures that all buffered spans are processed and ready for inspection.
171287
func (t *civisibilitymocktracer) Flush() {
172288
t.mock.Flush()
173-
t.real.Flush()
289+
t.realMu.RLock()
290+
defer t.realMu.RUnlock()
291+
if t.real != nil {
292+
t.real.Flush()
293+
}
174294
}

0 commit comments

Comments
 (0)