14
14
import static net .bytebuddy .matcher .ElementMatchers .named ;
15
15
import static net .bytebuddy .matcher .ElementMatchers .namedOneOf ;
16
16
17
+ import io .opentelemetry .api .trace .SpanKind ;
17
18
import io .opentelemetry .context .Context ;
18
19
import io .opentelemetry .context .Scope ;
19
20
import io .opentelemetry .instrumentation .api .annotation .support .async .AsyncOperationEndSupport ;
20
21
import io .opentelemetry .instrumentation .api .incubator .semconv .util .ClassAndMethod ;
21
22
import io .opentelemetry .javaagent .extension .instrumentation .TypeInstrumentation ;
22
23
import io .opentelemetry .javaagent .extension .instrumentation .TypeTransformer ;
24
+ import java .util .Locale ;
23
25
import java .util .Set ;
26
+ import java .util .logging .Level ;
27
+ import java .util .logging .Logger ;
28
+ import java .util .stream .Collectors ;
24
29
import net .bytebuddy .asm .Advice ;
25
30
import net .bytebuddy .description .type .TypeDescription ;
26
31
import net .bytebuddy .implementation .bytecode .assign .Assigner ;
27
32
import net .bytebuddy .matcher .ElementMatcher ;
28
33
29
34
public class MethodInstrumentation implements TypeInstrumentation {
30
35
private final String className ;
31
- private final Set <String > methodNames ;
36
+ private final Set <String > internalMethodNames ;
37
+ private final Set <String > serverMethodNames ;
38
+ private final Set <String > clientMethodNames ;
39
+
40
+ private static final Logger logger = Logger .getLogger (MethodInstrumentation .class .getName ());
32
41
33
42
public MethodInstrumentation (String className , Set <String > methodNames ) {
34
43
this .className = className ;
35
- this .methodNames = methodNames ;
44
+ this .internalMethodNames = filterMethodNames (className , methodNames , SpanKind .INTERNAL , true );
45
+ this .serverMethodNames = filterMethodNames (className , methodNames , SpanKind .CLIENT , false );
46
+ this .clientMethodNames = filterMethodNames (className , methodNames , SpanKind .SERVER , false );
47
+ }
48
+
49
+ private static Set <String > filterMethodNames (
50
+ String className , Set <String > methodNames , SpanKind kind , boolean isDefault ) {
51
+ String suffix = "=" + kind .name ();
52
+ Set <String > methods =
53
+ methodNames .stream ()
54
+ .filter (
55
+ methodName ->
56
+ methodName .toUpperCase (Locale .ROOT ).endsWith (suffix )
57
+ || (!methodName .contains ("=" ) && isDefault ))
58
+ .map (
59
+ methodName ->
60
+ methodName .contains ("=" )
61
+ ? methodName .substring (0 , methodName .indexOf ('=' ))
62
+ : methodName )
63
+ .collect (Collectors .toSet ());
64
+ logMethods (className , methods , kind );
65
+ return methods ;
66
+ }
67
+
68
+ private static void logMethods (String className , Set <String > methodNames , SpanKind spanKind ) {
69
+ if (!logger .isLoggable (Level .FINE ) || methodNames .isEmpty ()) {
70
+ return ;
71
+ }
72
+ logger .log (
73
+ Level .FINE ,
74
+ "Tracing class {0} with methods {1} using span kind {2}" ,
75
+ new Object [] {
76
+ className ,
77
+ methodNames .stream ()
78
+ .map (name -> className + "." + name )
79
+ .collect (Collectors .joining (", " )),
80
+ spanKind .name ()
81
+ });
36
82
}
37
83
38
84
@ Override
@@ -55,6 +101,13 @@ public ElementMatcher<TypeDescription> typeMatcher() {
55
101
56
102
@ Override
57
103
public void transform (TypeTransformer transformer ) {
104
+ applyMethodTransformation (transformer , internalMethodNames , "$InternalMethodAdvice" );
105
+ applyMethodTransformation (transformer , clientMethodNames , "$ClientMethodAdvice" );
106
+ applyMethodTransformation (transformer , serverMethodNames , "$ServerMethodAdvice" );
107
+ }
108
+
109
+ private static void applyMethodTransformation (
110
+ TypeTransformer transformer , Set <String > methodNames , String methodAdvice ) {
58
111
transformer .applyAdviceToMethod (
59
112
namedOneOf (methodNames .toArray (new String [0 ])).and (isMethod ()),
60
113
mapping ->
@@ -63,24 +116,109 @@ public void transform(TypeTransformer transformer) {
63
116
(instrumentedType , instrumentedMethod , assigner , argumentHandler , sort ) ->
64
117
Advice .OffsetMapping .Target .ForStackManipulation .of (
65
118
instrumentedMethod .getReturnType ().asErasure ())),
66
- MethodInstrumentation .class .getName () + "$MethodAdvice" );
119
+ MethodInstrumentation .class .getName () + methodAdvice );
67
120
}
68
121
69
122
// custom annotation that represents the return type of the method
70
123
@interface MethodReturnType {}
71
124
72
125
@ SuppressWarnings ("unused" )
73
- public static class MethodAdvice {
126
+ public static class InternalMethodAdvice {
74
127
75
128
@ Advice .OnMethodEnter (suppress = Throwable .class )
76
129
public static void onEnter (
77
130
@ Advice .Origin ("#t" ) Class <?> declaringClass ,
78
131
@ Advice .Origin ("#m" ) String methodName ,
79
- @ Advice .Local ("otelMethod" ) ClassAndMethod classAndMethod ,
132
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
80
133
@ Advice .Local ("otelContext" ) Context context ,
81
134
@ Advice .Local ("otelScope" ) Scope scope ) {
82
135
Context parentContext = currentContext ();
83
- classAndMethod = ClassAndMethod .create (declaringClass , methodName );
136
+ classAndMethod =
137
+ MethodAndType .create (
138
+ ClassAndMethod .create (declaringClass , methodName ), SpanKind .INTERNAL );
139
+
140
+ if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
141
+ return ;
142
+ }
143
+
144
+ context = instrumenter ().start (parentContext , classAndMethod );
145
+ scope = context .makeCurrent ();
146
+ }
147
+
148
+ @ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
149
+ public static void stopSpan (
150
+ @ MethodReturnType Class <?> methodReturnType ,
151
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
152
+ @ Advice .Local ("otelContext" ) Context context ,
153
+ @ Advice .Local ("otelScope" ) Scope scope ,
154
+ @ Advice .Return (typing = Assigner .Typing .DYNAMIC , readOnly = false ) Object returnValue ,
155
+ @ Advice .Thrown Throwable throwable ) {
156
+ if (scope == null ) {
157
+ return ;
158
+ }
159
+ scope .close ();
160
+
161
+ returnValue =
162
+ AsyncOperationEndSupport .create (instrumenter (), Void .class , methodReturnType )
163
+ .asyncEnd (context , classAndMethod , returnValue , throwable );
164
+ }
165
+ }
166
+
167
+ @ SuppressWarnings ("unused" )
168
+ public static class ClientMethodAdvice {
169
+
170
+ @ Advice .OnMethodEnter (suppress = Throwable .class )
171
+ public static void onEnter (
172
+ @ Advice .Origin ("#t" ) Class <?> declaringClass ,
173
+ @ Advice .Origin ("#m" ) String methodName ,
174
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
175
+ @ Advice .Local ("otelContext" ) Context context ,
176
+ @ Advice .Local ("otelScope" ) Scope scope ) {
177
+ Context parentContext = currentContext ();
178
+ classAndMethod =
179
+ MethodAndType .create (ClassAndMethod .create (declaringClass , methodName ), SpanKind .CLIENT );
180
+
181
+ if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
182
+ return ;
183
+ }
184
+
185
+ context = instrumenter ().start (parentContext , classAndMethod );
186
+ scope = context .makeCurrent ();
187
+ }
188
+
189
+ @ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
190
+ public static void stopSpan (
191
+ @ MethodReturnType Class <?> methodReturnType ,
192
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
193
+ @ Advice .Local ("otelContext" ) Context context ,
194
+ @ Advice .Local ("otelScope" ) Scope scope ,
195
+ @ Advice .Return (typing = Assigner .Typing .DYNAMIC , readOnly = false ) Object returnValue ,
196
+ @ Advice .Thrown Throwable throwable ) {
197
+ if (scope == null ) {
198
+ return ;
199
+ }
200
+ scope .close ();
201
+
202
+ returnValue =
203
+ AsyncOperationEndSupport .create (instrumenter (), Void .class , methodReturnType )
204
+ .asyncEnd (context , classAndMethod , returnValue , throwable );
205
+ }
206
+ }
207
+
208
+ @ SuppressWarnings ("unused" )
209
+ public static class ServerMethodAdvice {
210
+
211
+ @ Advice .OnMethodEnter (suppress = Throwable .class )
212
+ public static void onEnter (
213
+ @ Advice .Origin ("#t" ) Class <?> declaringClass ,
214
+ @ Advice .Origin ("#m" ) String methodName ,
215
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
216
+ @ Advice .Local ("otelContext" ) Context context ,
217
+ @ Advice .Local ("otelScope" ) Scope scope ) {
218
+ Context parentContext = currentContext ();
219
+ classAndMethod =
220
+ MethodAndType .create (ClassAndMethod .create (declaringClass , methodName ), SpanKind .SERVER );
221
+
84
222
if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
85
223
return ;
86
224
}
@@ -92,7 +230,7 @@ public static void onEnter(
92
230
@ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
93
231
public static void stopSpan (
94
232
@ MethodReturnType Class <?> methodReturnType ,
95
- @ Advice .Local ("otelMethod" ) ClassAndMethod classAndMethod ,
233
+ @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
96
234
@ Advice .Local ("otelContext" ) Context context ,
97
235
@ Advice .Local ("otelScope" ) Scope scope ,
98
236
@ Advice .Return (typing = Assigner .Typing .DYNAMIC , readOnly = false ) Object returnValue ,
0 commit comments