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