@@ -62,76 +62,7 @@ export class UIProfiler {
6262 this . _client = client ;
6363 this . _sessionSampled = sessionSampled ;
6464
65- client . on ( 'spanStart' , span => {
66- if ( ! this . _sessionSampled ) {
67- DEBUG_BUILD && debug . log ( '[Profiling] Session not sampled because of negative sampling decision.' ) ;
68- return ;
69- }
70- if ( span !== getRootSpan ( span ) ) {
71- return ;
72- }
73- // Only count sampled root spans
74- if ( ! span . isRecording ( ) ) {
75- DEBUG_BUILD && debug . log ( '[Profiling] Discarding profile because root span was not sampled.' ) ;
76- return ;
77- }
78-
79- // Matching root spans with profiles
80- getGlobalScope ( ) . setContext ( 'profile' , {
81- profiler_id : this . _profilerId ,
82- } ) ;
83-
84- const spanId = span . spanContext ( ) . spanId ;
85- if ( ! spanId ) {
86- return ;
87- }
88- if ( this . _activeRootSpanIds . has ( spanId ) ) {
89- return ;
90- }
91-
92- this . _activeRootSpanIds . add ( spanId ) ;
93- const rootSpanCount = this . _activeRootSpanIds . size ;
94-
95- const timeout = setTimeout ( ( ) => {
96- this . _onRootSpanTimeout ( spanId ) ;
97- } , MAX_ROOT_SPAN_PROFILE_MS ) ;
98- this . _rootSpanTimeouts . set ( spanId , timeout ) ;
99-
100- if ( rootSpanCount === 1 ) {
101- DEBUG_BUILD &&
102- debug . log (
103- `[Profiling] Root span with ID ${ spanId } started. Will continue profiling for as long as there are active root spans (currently: ${ rootSpanCount } ).` ,
104- ) ;
105-
106- this . start ( ) ;
107- }
108- } ) ;
109-
110- client . on ( 'spanEnd' , span => {
111- if ( ! this . _sessionSampled ) {
112- return ;
113- }
114-
115- const spanId = span . spanContext ( ) . spanId ;
116- if ( ! spanId || ! this . _activeRootSpanIds . has ( spanId ) ) {
117- return ;
118- }
119-
120- this . _activeRootSpanIds . delete ( spanId ) ;
121- const rootSpanCount = this . _activeRootSpanIds . size ;
122-
123- DEBUG_BUILD &&
124- debug . log (
125- `[Profiling] Root span with ID ${ spanId } ended. Will continue profiling for as long as there are active root spans (currently: ${ rootSpanCount } ).` ,
126- ) ;
127- if ( rootSpanCount === 0 ) {
128- this . _collectCurrentChunk ( ) . catch ( e => {
129- DEBUG_BUILD && debug . error ( '[Profiling] Failed to collect current profile chunk on `spanEnd`:' , e ) ;
130- } ) ;
131-
132- this . stop ( ) ;
133- }
134- } ) ;
65+ this . _setupTraceLifecycleListeners ( client ) ;
13566 }
13667
13768 /**
@@ -170,6 +101,9 @@ export class UIProfiler {
170101
171102 DEBUG_BUILD && debug . log ( '[Profiling] Started profiling with profile ID:' , this . _profilerId ) ;
172103
104+ // Expose profiler_id to match root spans with profiles
105+ getGlobalScope ( ) . setContext ( 'profile' , { profiler_id : this . _profilerId } ) ;
106+
173107 this . _startProfilerInstance ( ) ;
174108
175109 if ( ! this . _profiler ) {
@@ -203,6 +137,63 @@ export class UIProfiler {
203137 } ) ;
204138 }
205139
140+ /** Trace-mode: attach spanStart/spanEnd listeners. */
141+ private _setupTraceLifecycleListeners ( client : Client ) : void {
142+ client . on ( 'spanStart' , span => {
143+ if ( ! this . _sessionSampled ) {
144+ DEBUG_BUILD && debug . log ( '[Profiling] Session not sampled because of negative sampling decision.' ) ;
145+ return ;
146+ }
147+ if ( span !== getRootSpan ( span ) ) {
148+ return ; // only care about root spans
149+ }
150+ // Only count sampled root spans
151+ if ( ! span . isRecording ( ) ) {
152+ DEBUG_BUILD && debug . log ( '[Profiling] Discarding profile because root span was not sampled.' ) ;
153+ return ;
154+ }
155+
156+ const spanId = span . spanContext ( ) . spanId ;
157+ if ( ! spanId || this . _activeRootSpanIds . has ( spanId ) ) {
158+ return ;
159+ }
160+
161+ this . _registerTraceRootSpan ( spanId ) ;
162+
163+ const rootSpanCount = this . _activeRootSpanIds . size ;
164+ if ( rootSpanCount === 1 ) {
165+ DEBUG_BUILD &&
166+ debug . log (
167+ `[Profiling] Root span ${ spanId } started. Profiling active while there are active root spans (count=${ rootSpanCount } ).` ,
168+ ) ;
169+ this . start ( ) ;
170+ }
171+ } ) ;
172+
173+ client . on ( 'spanEnd' , span => {
174+ if ( ! this . _sessionSampled ) {
175+ return ;
176+ }
177+ const spanId = span . spanContext ( ) . spanId ;
178+ if ( ! spanId || ! this . _activeRootSpanIds . has ( spanId ) ) {
179+ return ;
180+ }
181+ this . _activeRootSpanIds . delete ( spanId ) ;
182+ const rootSpanCount = this . _activeRootSpanIds . size ;
183+
184+ DEBUG_BUILD &&
185+ debug . log (
186+ `[Profiling] Root span with ID ${ spanId } ended. Will continue profiling for as long as there are active root spans (currently: ${ rootSpanCount } ).` ,
187+ ) ;
188+ if ( rootSpanCount === 0 ) {
189+ this . _collectCurrentChunk ( ) . catch ( e => {
190+ DEBUG_BUILD && debug . error ( '[Profiling] Failed to collect current profile chunk on last `spanEnd`:' , e ) ;
191+ } ) ;
192+ this . stop ( ) ;
193+ }
194+ } ) ;
195+ }
196+
206197 /**
207198 * Resets profiling information from scope and resets running state
208199 */
@@ -219,6 +210,13 @@ export class UIProfiler {
219210 this . _rootSpanTimeouts . clear ( ) ;
220211 }
221212
213+ /** Register root span and schedule safeguard timeout (trace mode). */
214+ private _registerTraceRootSpan ( spanId : string ) : void {
215+ this . _activeRootSpanIds . add ( spanId ) ;
216+ const timeout = setTimeout ( ( ) => this . _onRootSpanTimeout ( spanId ) , MAX_ROOT_SPAN_PROFILE_MS ) ;
217+ this . _rootSpanTimeouts . set ( spanId , timeout ) ;
218+ }
219+
222220 /**
223221 * Start a profiler instance if needed.
224222 */
0 commit comments