2121import io .opentelemetry .instrumentation .api .incubator .semconv .util .ClassAndMethod ;
2222import io .opentelemetry .javaagent .extension .instrumentation .TypeInstrumentation ;
2323import io .opentelemetry .javaagent .extension .instrumentation .TypeTransformer ;
24- import java .util .Set ;
24+ import java .util .Map ;
2525import net .bytebuddy .asm .Advice ;
2626import net .bytebuddy .description .type .TypeDescription ;
2727import net .bytebuddy .implementation .bytecode .assign .Assigner ;
2828import net .bytebuddy .matcher .ElementMatcher ;
29+ import net .bytebuddy .utility .JavaConstant ;
2930
31+ @ SuppressWarnings ("EnumOrdinal" )
3032public class MethodInstrumentation implements TypeInstrumentation {
3133 private final String className ;
32- private final Set <String > internalMethodNames ;
33- private final Set <String > serverMethodNames ;
34- private final Set <String > clientMethodNames ;
35-
36- public MethodInstrumentation (
37- String className ,
38- Set <String > internalMethodNames ,
39- Set <String > serverMethodNames ,
40- Set <String > clientMethodNames ) {
34+ private final Map <String , SpanKind > methodNames ;
35+
36+ public MethodInstrumentation (String className , Map <String , SpanKind > methodNames ) {
4137 this .className = className ;
42- this .internalMethodNames = internalMethodNames ;
43- this .serverMethodNames = serverMethodNames ;
44- this .clientMethodNames = clientMethodNames ;
38+ this .methodNames = methodNames ;
4539 }
4640
4741 @ Override
@@ -64,123 +58,46 @@ public ElementMatcher<TypeDescription> typeMatcher() {
6458
6559 @ Override
6660 public void transform (TypeTransformer transformer ) {
67- applyMethodTransformation (transformer , internalMethodNames , "$InternalMethodAdvice" );
68- applyMethodTransformation (transformer , clientMethodNames , "$ClientMethodAdvice" );
69- applyMethodTransformation (transformer , serverMethodNames , "$ServerMethodAdvice" );
70- }
71-
72- private static void applyMethodTransformation (
73- TypeTransformer transformer , Set <String > methodNames , String methodAdvice ) {
7461 transformer .applyAdviceToMethod (
75- namedOneOf (methodNames .toArray (new String [0 ])).and (isMethod ()),
62+ namedOneOf (methodNames .keySet (). toArray (new String [0 ])).and (isMethod ()),
7663 mapping ->
77- mapping .bind (
78- MethodReturnType .class ,
79- (instrumentedType , instrumentedMethod , assigner , argumentHandler , sort ) ->
80- Advice .OffsetMapping .Target .ForStackManipulation .of (
81- instrumentedMethod .getReturnType ().asErasure ())),
82- MethodInstrumentation .class .getName () + methodAdvice );
64+ mapping
65+ .bind (
66+ MethodReturnType .class ,
67+ (instrumentedType , instrumentedMethod , assigner , argumentHandler , sort ) ->
68+ Advice .OffsetMapping .Target .ForStackManipulation .of (
69+ instrumentedMethod .getReturnType ().asErasure ()))
70+ .bind (
71+ SpanKindOrdinal .class ,
72+ (instrumentedType , instrumentedMethod , assigner , argumentHandler , sort ) ->
73+ Advice .OffsetMapping .Target .ForStackManipulation .of (
74+ JavaConstant .Simple .ofLoaded (
75+ methodNames .get (instrumentedMethod .getName ()).ordinal ()))),
76+ MethodInstrumentation .class .getName () + "$MethodAdvice" );
8377 }
8478
8579 // custom annotation that represents the return type of the method
8680 @interface MethodReturnType {}
8781
88- @ SuppressWarnings ("unused" )
89- public static class InternalMethodAdvice {
90-
91- @ Advice .OnMethodEnter (suppress = Throwable .class )
92- public static void onEnter (
93- @ Advice .Origin ("#t" ) Class <?> declaringClass ,
94- @ Advice .Origin ("#m" ) String methodName ,
95- @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
96- @ Advice .Local ("otelContext" ) Context context ,
97- @ Advice .Local ("otelScope" ) Scope scope ) {
98- Context parentContext = currentContext ();
99- classAndMethod =
100- MethodAndType .create (
101- ClassAndMethod .create (declaringClass , methodName ), SpanKind .INTERNAL );
102-
103- if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
104- return ;
105- }
106-
107- context = instrumenter ().start (parentContext , classAndMethod );
108- scope = context .makeCurrent ();
109- }
110-
111- @ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
112- public static void stopSpan (
113- @ MethodReturnType Class <?> methodReturnType ,
114- @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
115- @ Advice .Local ("otelContext" ) Context context ,
116- @ Advice .Local ("otelScope" ) Scope scope ,
117- @ Advice .Return (typing = Assigner .Typing .DYNAMIC , readOnly = false ) Object returnValue ,
118- @ Advice .Thrown Throwable throwable ) {
119- if (scope == null ) {
120- return ;
121- }
122- scope .close ();
123-
124- returnValue =
125- AsyncOperationEndSupport .create (instrumenter (), Void .class , methodReturnType )
126- .asyncEnd (context , classAndMethod , returnValue , throwable );
127- }
128- }
129-
130- @ SuppressWarnings ("unused" )
131- public static class ClientMethodAdvice {
132-
133- @ Advice .OnMethodEnter (suppress = Throwable .class )
134- public static void onEnter (
135- @ Advice .Origin ("#t" ) Class <?> declaringClass ,
136- @ Advice .Origin ("#m" ) String methodName ,
137- @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
138- @ Advice .Local ("otelContext" ) Context context ,
139- @ Advice .Local ("otelScope" ) Scope scope ) {
140- Context parentContext = currentContext ();
141- classAndMethod =
142- MethodAndType .create (ClassAndMethod .create (declaringClass , methodName ), SpanKind .CLIENT );
143-
144- if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
145- return ;
146- }
147-
148- context = instrumenter ().start (parentContext , classAndMethod );
149- scope = context .makeCurrent ();
150- }
151-
152- @ Advice .OnMethodExit (onThrowable = Throwable .class , suppress = Throwable .class )
153- public static void stopSpan (
154- @ MethodReturnType Class <?> methodReturnType ,
155- @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
156- @ Advice .Local ("otelContext" ) Context context ,
157- @ Advice .Local ("otelScope" ) Scope scope ,
158- @ Advice .Return (typing = Assigner .Typing .DYNAMIC , readOnly = false ) Object returnValue ,
159- @ Advice .Thrown Throwable throwable ) {
160- if (scope == null ) {
161- return ;
162- }
163- scope .close ();
164-
165- returnValue =
166- AsyncOperationEndSupport .create (instrumenter (), Void .class , methodReturnType )
167- .asyncEnd (context , classAndMethod , returnValue , throwable );
168- }
169- }
82+ // custom annotation that represents the SpanKind of the method
83+ @interface SpanKindOrdinal {}
17084
17185 @ SuppressWarnings ("unused" )
172- public static class ServerMethodAdvice {
86+ public static class MethodAdvice {
17387
17488 @ Advice .OnMethodEnter (suppress = Throwable .class )
17589 public static void onEnter (
90+ @ SpanKindOrdinal int spanKindOrdinal ,
17691 @ Advice .Origin ("#t" ) Class <?> declaringClass ,
17792 @ Advice .Origin ("#m" ) String methodName ,
17893 @ Advice .Local ("otelMethod" ) MethodAndType classAndMethod ,
17994 @ Advice .Local ("otelContext" ) Context context ,
18095 @ Advice .Local ("otelScope" ) Scope scope ) {
18196 Context parentContext = currentContext ();
18297 classAndMethod =
183- MethodAndType .create (ClassAndMethod .create (declaringClass , methodName ), SpanKind .SERVER );
98+ MethodAndType .create (
99+ ClassAndMethod .create (declaringClass , methodName ),
100+ SpanKind .values ()[spanKindOrdinal ]);
184101
185102 if (!instrumenter ().shouldStart (parentContext , classAndMethod )) {
186103 return ;
0 commit comments