6
6
import datadog .trace .core .PendingTrace ;
7
7
import java .io .IOException ;
8
8
import java .lang .ref .WeakReference ;
9
+ import java .util .HashMap ;
9
10
import java .util .HashSet ;
11
+ import java .util .Map ;
10
12
import java .util .Set ;
11
13
import java .util .zip .ZipOutputStream ;
12
14
@@ -19,31 +21,38 @@ public static void register() {
19
21
}
20
22
21
23
private static final Set <WeakReference <DDSpan >> rootSpans = new HashSet <>();
24
+ private static final Map <String , WeakReference <DDSpan >> idToWeakRef = new HashMap <>();
22
25
23
- public static void addActiveSpan (final AgentSpan rootSpan ) {
26
+ public static void addUnfinishedRootSpan (final AgentSpan rootSpan ) {
24
27
DDSpan rootDDSpan = (DDSpan ) rootSpan ;
25
28
WeakReference <DDSpan > weakRootSpan = new WeakReference <>(rootDDSpan );
29
+ idToWeakRef .put (getSpanIdString (rootSpan ), weakRootSpan );
26
30
rootSpans .add (weakRootSpan );
27
31
}
28
32
29
33
public static void suspendRootSpan (final AgentSpan rootSpan ) {
30
- DDSpan rootDDSpan = (DDSpan ) rootSpan ;
31
- WeakReference <DDSpan > weakRootSpan = new WeakReference <>(rootDDSpan );
34
+ WeakReference <DDSpan > weakRootSpan = idToWeakRef .get (getSpanIdString (rootSpan ));
32
35
rootSpans .remove (weakRootSpan );
33
36
}
34
37
35
38
@ Override
36
39
public void addReportToFlare (ZipOutputStream zip ) throws IOException {
40
+ StringBuilder dumpText = new StringBuilder ();
37
41
for (WeakReference <DDSpan > weakRootSpan : rootSpans ) {
38
42
DDSpan rootSpan = weakRootSpan .get ();
39
43
if (rootSpan != null ) {
40
44
PendingTrace trace = (PendingTrace ) rootSpan .context ().getTraceCollector ();
41
45
if (trace != null ) {
42
46
for (DDSpan span : trace .getSpans ()) {
43
- TracerFlare . addText ( zip , "trace_dump.txt" , span .toString ());
47
+ dumpText . append ( span .toString ()). append ( " \n " );
44
48
}
45
49
}
46
50
}
47
51
}
52
+ TracerFlare .addText (zip , "trace_dump.txt" , dumpText .toString ());
53
+ }
54
+
55
+ private static String getSpanIdString (final AgentSpan rootSpan ) {
56
+ return String .valueOf (rootSpan .getSpanId ());
48
57
}
49
58
}
0 commit comments