How properly ThreadLocal Context in the ByteBuddy Instrumentation Advice

Question

Hi I'm trying to build a trace agent where app agent propagates the origin application name across the multiple network calls using the HTTPUrlConnection.

So far, I was able to invoke the Advice correctly with the new value.

however, when ApplicationTraceContext updated with the new value from the application code, it does not reflect with in the advice.

Ex.

• First invocation APP Name: *ClientAPP * -> Advice Reads *ClientAPP *
• App updates the context to *BatchAPP * -> Advice Still Reads *ClientAPP *

Agent

new AgentBuilder.Default()
  //.with(new MyListener())
  .ignore(none())
  .type(typeMatcher())
  .transform((builder, typeDescription, classLoader, module) -> {
    System.out.println("Transforming inside: " + typeDescription.getName());
    return builder
      .visit(
        Advice.withCustomMapping()
          .bind(AgentArguments.class, ApplicationTraceContext.getRequestOriginApplicationName())
          .to(HttpUrlConnectionAdvice.class)
          .on(
            namedOneOf("connect", "getOutputStream", "getInputStream", "getResponseCode")
          ));

  })
  .installOn(inst);

Advice

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void methodEnter(
  @Advice.This HttpURLConnection connection,
  @Advice.FieldValue("connected") boolean connected,
  @AgentArguments Object agentArguments
)
{
  System.out.println("Intercepted Connection: " + connection.getClass().getName());

  System.out.println("Agent Args from context" + agentArguments);
  if (agentArguments != null && !agentArguments.toString().isEmpty()) {
    System.out.println("Setting header: X-Application-Name from agent arguments " + agentArguments);
    connection.setRequestProperty(Constants.ORIGIN_APPLICATION_NAME_HEADER_KEY, agentArguments.toString());
  }
  System.out.println("Request Headers: " + connection.getRequestProperties());
}

Alos I tried directly accessing the ApplicationTraceContext.getRequestOriginApplicationName() within the Advice, it does not return the value or stopes executing the remaining advice.

Any suggestions to resolve this is appreciated.

Additionally, tried with the custom offset biding, the result is still same.

Advice.withCustomMapping()
  .bind(AgentArguments.class, new AgentOffsetMapping())
  .to(HttpUrlConnectionAdvice.class)

public class AgentOffsetMapping implements Advice.OffsetMapping {
  @Override
  public Target resolve(TypeDescription instrumentedType, MethodDescription instrumentedMethod, Assigner assigner, Advice.ArgumentHandler argumentHandler, Sort sort) {
    return new Target() {
      @Override
      public StackManipulation resolveRead() {
        try {
          MethodDescription.ForLoadedMethod getThreadLocalValue = new MethodDescription.ForLoadedMethod(
            ThreadLocalHolder.class.getMethod("getThreadLocalValue")
          );
          return MethodInvocation.invoke(
            getThreadLocalValue);
        }
        catch (NoSuchMethodException e) {
          System.out.println("Failed to resolve read" + e.getMessage());
          throw new RuntimeException(e);
        }
      }

      @Override
      public StackManipulation resolveWrite() {
        throw new UnsupportedOperationException("Cannot write to @MyAnnotation parameter");
      }

      @Override
      public StackManipulation resolveIncrement(int value) {
        throw new UnsupportedOperationException("Cannot write to @MyAnnotation parameter");
      }
    };
  }
}

Answer 1

The code in the advice method is inlined into the target method. If the target method's class cannot see the invoked code, it will not be invokable, but trigger a NoClassDefFoundError. Since you suppress all Throwables, the error does not surface.

You would likely need to inject the class you want to invoke into the boot loader. Normally, this will be some type of API class with a static field. The boot loader is visible from all class loaders, also from your agent. You can now set the field, and invoke the API methods from your advice on the instance that was set.