Generate InstrumentationModule#contextStore() with muzzle (#2775)

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorInstrumentationModule.java

@@ -8,14 +8,9 @@ package io.opentelemetry.javaagent.instrumentation.akkaactor;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.Callable;
 
 @AutoService(InstrumentationModule.class)
 public class AkkaActorInstrumentationModule extends InstrumentationModule {
@@ -32,16 +27,6 @@ public class AkkaActorInstrumentationModule extends InstrumentationModule {
         new AkkaActorCellInstrumentation());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put(Runnable.class.getName(), State.class.getName());
-    map.put(Callable.class.getName(), State.class.getName());
-    map.put(AkkaForkJoinTaskInstrumentation.TASK_CLASS_NAME, State.class.getName());
-    map.put("akka.dispatch.Envelope", State.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
-
   @Override
   protected boolean defaultEnabled() {
     return false;

instrumentation/async-http-client/async-http-client-1.9/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/AsyncHttpClientInstrumentationModule.java

@@ -6,13 +6,10 @@
 package io.opentelemetry.javaagent.instrumentation.asynchttpclient;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.javaagent.instrumentation.api.Pair;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(io.opentelemetry.javaagent.tooling.InstrumentationModule.class)
 public class AsyncHttpClientInstrumentationModule
@@ -25,9 +22,4 @@ public class AsyncHttpClientInstrumentationModule
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new RequestInstrumentation(), new ResponseInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("com.ning.http.client.AsyncHandler", Pair.class.getName());
-  }
 }

instrumentation/async-http-client/async-http-client-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/asynchttpclient/AsyncHttpClientInstrumentationModule.java

@@ -8,11 +8,8 @@ package io.opentelemetry.javaagent.instrumentation.asynchttpclient;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.context.Context;
-import io.opentelemetry.javaagent.instrumentation.api.Pair;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(io.opentelemetry.javaagent.tooling.InstrumentationModule.class)
 public class AsyncHttpClientInstrumentationModule
@@ -28,12 +25,4 @@ public class AsyncHttpClientInstrumentationModule
         new ResponseInstrumentation(),
         new RequestSenderInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> stores = new java.util.HashMap<>();
-    stores.put("org.asynchttpclient.AsyncHandler", Pair.class.getName());
-    stores.put("org.asynchttpclient.Request", Context.class.getName());
-    return stores;
-  }
 }

instrumentation/couchbase/couchbase-2.6/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/couchbase/v2_6/CouchbaseInstrumentationModule.java

@@ -6,14 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.couchbase.v2_6;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class CouchbaseInstrumentationModule extends InstrumentationModule {
@@ -25,9 +22,4 @@ public class CouchbaseInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new CouchbaseCoreInstrumentation(), new CouchbaseNetworkInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("com.couchbase.client.core.message.CouchbaseRequest", Span.class.getName());
-  }
 }

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/javaconcurrent/ExecutorInstrumentationModule.java

@@ -8,16 +8,9 @@ package io.opentelemetry.javaagent.instrumentation.javaconcurrent;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.Callable;
-import java.util.concurrent.ForkJoinTask;
-import java.util.concurrent.Future;
 
 @AutoService(InstrumentationModule.class)
 public class ExecutorInstrumentationModule extends InstrumentationModule {
@@ -34,14 +27,4 @@ public class ExecutorInstrumentationModule extends InstrumentationModule {
         new JavaForkJoinTaskInstrumentation(),
         new RunnableInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put(Callable.class.getName(), State.class.getName());
-    map.put(ForkJoinTask.class.getName(), State.class.getName());
-    map.put(Future.class.getName(), State.class.getName());
-    map.put(Runnable.class.getName(), State.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
 }

instrumentation/google-http-client-1.19/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/googlehttpclient/GoogleHttpClientInstrumentationModule.java

@@ -8,7 +8,6 @@ package io.opentelemetry.javaagent.instrumentation.googlehttpclient;
 import static io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.instrumentation.googlehttpclient.GoogleHttpClientTracer.tracer;
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -38,11 +37,6 @@ public class GoogleHttpClientInstrumentationModule extends InstrumentationModule
     super("google-http-client", "google-http-client-1.19");
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("com.google.api.client.http.HttpRequest", Context.class.getName());
-  }
-
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new HttpRequestInstrumentation());

instrumentation/guava-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/guava/GuavaInstrumentationModule.java

@@ -41,11 +41,6 @@ public class GuavaInstrumentationModule extends InstrumentationModule {
     return singletonList(new ListenableFutureInstrumentation());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(Runnable.class.getName(), State.class.getName());
-  }
-
   public static class ListenableFutureInstrumentation implements TypeInstrumentation {
     @Override
     public ElementMatcher<TypeDescription> typeMatcher() {

instrumentation/hibernate/hibernate-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/v3_3/HibernateInstrumentationModule.java

@@ -9,13 +9,9 @@ import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMa
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -43,15 +39,4 @@ public class HibernateInstrumentationModule extends InstrumentationModule {
         new SessionInstrumentation(),
         new TransactionInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put("org.hibernate.Criteria", Context.class.getName());
-    map.put("org.hibernate.Query", Context.class.getName());
-    map.put("org.hibernate.Session", Context.class.getName());
-    map.put("org.hibernate.StatelessSession", Context.class.getName());
-    map.put("org.hibernate.Transaction", Context.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
 }

instrumentation/hibernate/hibernate-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/v4_0/HibernateInstrumentationModule.java

@@ -8,13 +8,9 @@ package io.opentelemetry.javaagent.instrumentation.hibernate.v4_0;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class HibernateInstrumentationModule extends InstrumentationModule {
@@ -32,14 +28,4 @@ public class HibernateInstrumentationModule extends InstrumentationModule {
         new SessionInstrumentation(),
         new TransactionInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put("org.hibernate.Criteria", Context.class.getName());
-    map.put("org.hibernate.Query", Context.class.getName());
-    map.put("org.hibernate.SharedSessionContract", Context.class.getName());
-    map.put("org.hibernate.Transaction", Context.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
 }

instrumentation/hibernate/hibernate-4.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/v4_3/HibernateInstrumentationModule.java

@@ -8,13 +8,9 @@ package io.opentelemetry.javaagent.instrumentation.hibernate.v4_3;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class HibernateInstrumentationModule extends InstrumentationModule {
@@ -26,12 +22,4 @@ public class HibernateInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new ProcedureCallInstrumentation(), new SessionInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put("org.hibernate.SharedSessionContract", Context.class.getName());
-    map.put("org.hibernate.procedure.ProcedureCall", Context.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
 }

instrumentation/http-url-connection/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpurlconnection/HttpUrlConnectionInstrumentationModule.java

@@ -10,7 +10,6 @@ import static io.opentelemetry.javaagent.instrumentation.httpurlconnection.HttpU
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.extendsClass;
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
@@ -52,11 +51,6 @@ public class HttpUrlConnectionInstrumentationModule extends InstrumentationModul
     return singletonList(new HttpUrlConnectionInstrumentation());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("java.net.HttpURLConnection", getClass().getName() + "$HttpUrlState");
-  }
-
   public static class HttpUrlConnectionInstrumentation implements TypeInstrumentation {
     @Override
     public ElementMatcher<TypeDescription> typeMatcher() {

instrumentation/jaxrs/jaxrs-2.0/jaxrs-2.0-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jaxrs/v2_0/JaxRsInstrumentationModule.java

@@ -7,14 +7,11 @@ package io.opentelemetry.javaagent.instrumentation.jaxrs.v2_0;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -37,9 +34,4 @@ public class JaxRsInstrumentationModule extends InstrumentationModule {
         new JaxRsAnnotationsInstrumentation(),
         new JaxRsAsyncResponseInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("javax.ws.rs.container.AsyncResponse", Context.class.getName());
-  }
 }

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/Jetty8InstrumentationModule.java

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.jetty.v8_0;
 
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -45,11 +44,6 @@ public class Jetty8InstrumentationModule extends InstrumentationModule {
         new JettyQueuedThreadPoolInstrumentation());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(Runnable.class.getName(), State.class.getName());
-  }
-
   public static class JettyQueuedThreadPoolInstrumentation implements TypeInstrumentation {
 
     @Override

instrumentation/jms-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jms/JmsInstrumentationModule.java

@@ -6,13 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.jms;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class JmsInstrumentationModule extends InstrumentationModule {
@@ -27,9 +25,4 @@ public class JmsInstrumentationModule extends InstrumentationModule {
         new JmsMessageListenerInstrumentation(),
         new JmsMessageProducerInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("javax.jms.MessageConsumer", MessageDestination.class.getName());
-  }
 }

instrumentation/log4j/log4j-1.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/v1_2/Log4j1InstrumentationModule.java

@@ -6,14 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.log4j.v1_2;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class Log4j1InstrumentationModule extends InstrumentationModule {
@@ -25,9 +22,4 @@ public class Log4j1InstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new CategoryInstrumentation(), new LoggingEventInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("org.apache.log4j.spi.LoggingEvent", Span.class.getName());
-  }
 }

instrumentation/logback/logback-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/logback/v1_0/LogbackInstrumentationModule.java

@@ -6,14 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.logback.v1_0;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class LogbackInstrumentationModule extends InstrumentationModule {
@@ -25,9 +22,4 @@ public class LogbackInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new LoggerInstrumentation(), new LoggingEventInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("ch.qos.logback.classic.spi.ILoggingEvent", Span.class.getName());
-  }
 }

instrumentation/netty/netty-3.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v3_8/NettyInstrumentationModule.java

@@ -10,9 +10,7 @@ import static java.util.Arrays.asList;
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class NettyInstrumentationModule extends InstrumentationModule {
@@ -28,10 +26,4 @@ public class NettyInstrumentationModule extends InstrumentationModule {
         new NettyChannelPipelineInstrumentation(),
         new DefaultChannelPipelineInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return Collections.singletonMap(
-        "org.jboss.netty.channel.Channel", ChannelTraceContext.class.getName());
-  }
 }

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/RmiContextPropagationInstrumentationModule.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.instrumentation.rmi.context;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.instrumentation.rmi.context.client.RmiClientContextInstrumentation;
@@ -14,7 +13,6 @@ import io.opentelemetry.javaagent.instrumentation.rmi.context.server.RmiServerCo
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class RmiContextPropagationInstrumentationModule extends InstrumentationModule {
@@ -26,10 +24,4 @@ public class RmiContextPropagationInstrumentationModule extends InstrumentationM
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new RmiClientContextInstrumentation(), new RmiServerContextInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    // caching if a connection can support enhanced format
-    return singletonMap("sun.rmi.transport.Connection", "java.lang.Boolean");
-  }
 }

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java

@@ -78,6 +78,7 @@ public class RmiClientContextInstrumentation implements TypeInstrumentation {
         return;
       }
 
+      // caching if a connection can support enhanced format
       ContextStore<Connection, Boolean> knownConnections =
           InstrumentationContext.get(Connection.class, Boolean.class);
 

instrumentation/scala-executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/scalaexecutors/ScalaConcurrentInstrumentationModule.java

@@ -5,18 +5,12 @@
 
 package io.opentelemetry.javaagent.instrumentation.scalaexecutors;
 
-import static io.opentelemetry.javaagent.instrumentation.scalaexecutors.ScalaForkJoinTaskInstrumentation.TASK_CLASS_NAME;
 import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.Callable;
 
 @AutoService(InstrumentationModule.class)
 public class ScalaConcurrentInstrumentationModule extends InstrumentationModule {
@@ -28,13 +22,4 @@ public class ScalaConcurrentInstrumentationModule extends InstrumentationModule
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new ScalaForkJoinPoolInstrumentation(), new ScalaForkJoinTaskInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> map = new HashMap<>();
-    map.put(Runnable.class.getName(), State.class.getName());
-    map.put(Callable.class.getName(), State.class.getName());
-    map.put(TASK_CLASS_NAME, State.class.getName());
-    return Collections.unmodifiableMap(map);
-  }
 }

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2InstrumentationModule.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.instrumentation.servlet.v2_2;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import com.google.auto.service.AutoService;
@@ -15,7 +14,6 @@ import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -38,9 +36,4 @@ public class Servlet2InstrumentationModule extends InstrumentationModule {
             "javax.servlet",
             Servlet2InstrumentationModule.class.getPackage().getName() + ".Servlet2Advice"));
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("javax.servlet.ServletResponse", Integer.class.getName());
-  }
 }

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/JakartaServletInstrumentationModule.java

@@ -5,8 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0;
 
-import static java.util.Collections.singletonMap;
-
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.async.AsyncContextInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.dispatcher.RequestDispatcherInstrumentation;
@@ -16,7 +14,6 @@ import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class JakartaServletInstrumentationModule extends InstrumentationModule {
@@ -39,11 +36,6 @@ public class JakartaServletInstrumentationModule extends InstrumentationModule {
             BASE_PACKAGE, adviceClassName(".dispatcher.RequestDispatcherAdvice")));
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(BASE_PACKAGE + ".RequestDispatcher", String.class.getName());
-  }
-
   private static String adviceClassName(String suffix) {
     return JakartaServletInstrumentationModule.class.getPackage().getName() + suffix;
   }

instrumentation/servlet/servlet-javax-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/javax/JavaxServletInstrumentationModule.java

@@ -5,8 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.javax;
 
-import static java.util.Collections.singletonMap;
-
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.dispatcher.RequestDispatcherInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.response.HttpServletResponseInstrumentation;
@@ -14,7 +12,6 @@ import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class JavaxServletInstrumentationModule extends InstrumentationModule {
@@ -33,11 +30,6 @@ public class JavaxServletInstrumentationModule extends InstrumentationModule {
             BASE_PACKAGE, adviceClassName(".dispatcher.RequestDispatcherAdvice")));
   }
 
-  @Override
-  protected Map<String, String> contextStore() {
-    return singletonMap(BASE_PACKAGE + ".RequestDispatcher", String.class.getName());
-  }
-
   private static String adviceClassName(String suffix) {
     return JavaxServletInstrumentationModule.class.getPackage().getName() + suffix;
   }

instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/SpringBatchInstrumentationModule.java

@@ -21,9 +21,7 @@ import io.opentelemetry.javaagent.instrumentation.spring.batch.step.StepBuilderH
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -37,17 +35,6 @@ public class SpringBatchInstrumentationModule extends InstrumentationModule {
     return hasClassesNamed("org.springframework.batch.core.jsr.launch.JsrJobOperator");
   }
 
-  @Override
-  protected Map<String, String> contextStore() {
-    Map<String, String> context = new HashMap<>();
-    String contextAndScope =
-        "io.opentelemetry.javaagent.instrumentation.spring.batch.ContextAndScope";
-    context.put("org.springframework.batch.core.JobExecution", contextAndScope);
-    context.put("org.springframework.batch.core.StepExecution", contextAndScope);
-    context.put("org.springframework.batch.core.scope.context.ChunkContext", contextAndScope);
-    return context;
-  }
-
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return Arrays.asList(

instrumentation/spring/spring-core-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/core/SpringCoreInstrumentationModule.java

@@ -7,14 +7,11 @@ package io.opentelemetry.javaagent.instrumentation.spring.core;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -28,11 +25,6 @@ public class SpringCoreInstrumentationModule extends InstrumentationModule {
     return hasClassesNamed("org.springframework.core.task.SimpleAsyncTaskExecutor");
   }
 
-  @Override
-  protected Map<String, String> contextStore() {
-    return singletonMap(Runnable.class.getName(), State.class.getName());
-  }
-
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new SimpleAsyncTaskExecutorInstrumentation());

instrumentation/vertx-web-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/client/VertxClientInstrumentationModule.java

@@ -6,13 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.vertx.client;
 
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.List;
-import java.util.Map;
 
 @AutoService(InstrumentationModule.class)
 public class VertxClientInstrumentationModule extends InstrumentationModule {
@@ -25,9 +23,4 @@ public class VertxClientInstrumentationModule extends InstrumentationModule {
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new HttpRequestInstrumentation());
   }
-
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap("io.vertx.core.http.HttpClientRequest", Contexts.class.getName());
-  }
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/InstrumentationModule.java

@@ -192,7 +192,7 @@ public abstract class InstrumentationModule {
   }
 
   private InstrumentationContextProvider getContextProvider() {
-    Map<String, String> contextStore = contextStore();
+    Map<String, String> contextStore = getMuzzleContextStoreClasses();
     if (!contextStore.isEmpty()) {
       return new FieldBackedProvider(getClass(), contextStore);
     } else {
@@ -303,6 +303,20 @@ public abstract class InstrumentationModule {
     return EMPTY;
   }
 
+  /**
+   * Returns a map of {@code class-name to context-class-name}. Keys (and their subclasses) will be
+   * associated with a context class stored in the value.
+   *
+   * <p>The actual implementation of this method is generated automatically during compilation by
+   * the {@link io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCodeGenerationPlugin}
+   * ByteBuddy plugin.
+   *
+   * <p><b>This method is generated automatically, do not override it.</b>
+   */
+  protected Map<String, String> getMuzzleContextStoreClasses() {
+    return Collections.emptyMap();
+  }
+
   /**
    * Instrumentation modules can override this method to provide additional helper classes that are
    * not located in instrumentation packages described in {@link InstrumentationClassPredicate} and
@@ -346,16 +360,6 @@ public abstract class InstrumentationModule {
   /** Returns a list of all individual type instrumentation in this module. */
   public abstract List<TypeInstrumentation> typeInstrumentations();
 
-  /**
-   * Context stores to define for this instrumentation.
-   *
-   * <p>A map of {@code class-name to context-class-name}. Keys (and their subclasses) will be
-   * associated with a context of the value.
-   */
-  protected Map<String, String> contextStore() {
-    return Collections.emptyMap();
-  }
-
   /**
    * Allows instrumentation modules to disable themselves by default, or to additionally disable
    * themselves on some other condition.

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/MuzzleCodeGenerator.java

@@ -10,6 +10,7 @@ import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference;
 import io.opentelemetry.javaagent.tooling.muzzle.matcher.ReferenceMatcher;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -42,6 +43,8 @@ class MuzzleCodeGenerator implements AsmVisitorWrapper {
   public static final String MUZZLE_REF_MATCHER_FIELD_NAME = "muzzleReferenceMatcher";
   public static final String MUZZLE_REF_MATCHER_METHOD_NAME = "getMuzzleReferenceMatcher";
   public static final String MUZZLE_HELPER_CLASSES_METHOD_NAME = "getMuzzleHelperClassNames";
+  public static final String MUZZLE_CONTEXT_STORE_CLASSES_METHOD_NAME =
+      "getMuzzleContextStoreClasses";
 
   @Override
   public int mergeWriter(int flags) {
@@ -138,6 +141,7 @@ class MuzzleCodeGenerator implements AsmVisitorWrapper {
       generateMuzzleHelperClassNamesMethod(collector);
       generateMuzzleReferenceMatcherMethod(collector);
       generateMuzzleReferenceMatcherField();
+      generateMuzzleContextStoreClassesMethod(collector);
       super.visitEnd();
     }
 
@@ -518,6 +522,63 @@ class MuzzleCodeGenerator implements AsmVisitorWrapper {
           null);
     }
 
+    private void generateMuzzleContextStoreClassesMethod(ReferenceCollector collector) {
+      /*
+       * protected Map<String, String> getMuzzleContextStoreClasses() {
+       *   Map<String, String> contextStore = new HashMap();
+       *   contextStore.put(..., ...);
+       *   return contextStore;
+       * }
+       */
+      MethodVisitor mv =
+          super.visitMethod(
+              Opcodes.ACC_PROTECTED,
+              MUZZLE_CONTEXT_STORE_CLASSES_METHOD_NAME,
+              "()Ljava/util/Map;",
+              null,
+              null);
+      mv.visitCode();
+
+      Map<String, String> contextStoreClasses = collector.getContextStoreClasses();
+
+      mv.visitTypeInsn(Opcodes.NEW, "java/util/HashMap");
+      // stack: map
+      mv.visitInsn(Opcodes.DUP);
+      // stack: map, map
+      mv.visitLdcInsn(contextStoreClasses.size());
+      // stack: map, map, size
+      mv.visitMethodInsn(Opcodes.INVOKESPECIAL, "java/util/HashMap", "<init>", "(I)V", false);
+      // stack: map
+      mv.visitVarInsn(Opcodes.ASTORE, 1);
+      // stack: <empty>
+
+      contextStoreClasses.forEach(
+          (className, contextClassName) -> {
+            mv.visitVarInsn(Opcodes.ALOAD, 1);
+            // stack: map
+            mv.visitLdcInsn(className);
+            // stack: map, className
+            mv.visitLdcInsn(contextClassName);
+            // stack: map, className, contextClassName
+            mv.visitMethodInsn(
+                Opcodes.INVOKEINTERFACE,
+                "java/util/Map",
+                "put",
+                "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
+                true);
+            // stack: previousValue
+            mv.visitInsn(Opcodes.POP);
+            // stack: <empty>
+          });
+
+      mv.visitVarInsn(Opcodes.ALOAD, 1);
+      // stack: map
+      mv.visitInsn(Opcodes.ARETURN);
+
+      mv.visitMaxs(0, 0);
+      mv.visitEnd();
+    }
+
     private static final Pattern ANONYMOUS_ENUM_CONSTANT_CLASS =
         Pattern.compile("(?<enumClass>.*)\\$[0-9]+$");
 

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/MuzzleCompilationException.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.muzzle.collector;
+
+public class MuzzleCompilationException extends RuntimeException {
+  public MuzzleCompilationException(String message) {
+    super(message);
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/ReferenceCollectingClassVisitor.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.tooling.muzzle.collector;
 
+import com.google.common.collect.EvictingQueue;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference;
@@ -115,6 +116,7 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
   // helper super classes which are themselves also helpers
   // this is needed for injecting the helper classes into the class loader in the correct order
   private final Set<String> helperSuperClasses = new HashSet<>();
+  private final Map<String, String> contextStoreClasses = new LinkedHashMap<>();
   private String refSourceClassName;
   private Type refSourceType;
 
@@ -137,6 +139,10 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
     return helperSuperClasses;
   }
 
+  Map<String, String> getContextStoreClasses() {
+    return contextStoreClasses;
+  }
+
   private void addExtendsReference(Reference ref) {
     addReference(ref);
     if (instrumentationClassPredicate.isInstrumentationClass(ref.getClassName())) {
@@ -240,7 +246,8 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
     // Additional references we could check
     // - Classes in signature (return type, params) and visible from this package
     return new AdviceReferenceMethodVisitor(
-        super.visitMethod(access, name, descriptor, signature, exceptions));
+        new InstrumentationContextMethodVisitor(
+            super.visitMethod(access, name, descriptor, signature, exceptions)));
   }
 
   private static VisibilityFlag computeVisibilityFlag(int access) {
@@ -332,6 +339,7 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
                 .withFlag(computeMinimumClassAccess(refSourceType, underlyingFieldType))
                 .build());
       }
+
       super.visitFieldInsn(opcode, owner, name, descriptor);
     }
 
@@ -349,6 +357,11 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
       //   * return type
       Type methodType = Type.getMethodType(descriptor);
 
+      Type ownerType =
+          owner.startsWith("[")
+              ? underlyingType(Type.getType(owner))
+              : Type.getType("L" + owner + ";");
+
       { // ref for method return type
         Type returnType = underlyingType(methodType.getReturnType());
         if (returnType.getSort() == Type.OBJECT) {
@@ -371,11 +384,6 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
         }
       }
 
-      Type ownerType =
-          owner.startsWith("[")
-              ? underlyingType(Type.getType(owner))
-              : Type.getType("L" + owner + ";");
-
       List<Reference.Flag> methodFlags = new ArrayList<>();
       methodFlags.add(
           opcode == Opcodes.INVOKESTATIC ? OwnershipFlag.STATIC : OwnershipFlag.NON_STATIC);
@@ -395,6 +403,7 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
                   methodType.getReturnType(),
                   methodType.getArgumentTypes())
               .build());
+
       super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
     }
 
@@ -408,6 +417,7 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
                 .withFlag(computeMinimumClassAccess(refSourceType, typeObj))
                 .build());
       }
+
       super.visitTypeInsn(opcode, type);
     }
 
@@ -456,4 +466,106 @@ class ReferenceCollectingClassVisitor extends ClassVisitor {
       super.visitLdcInsn(value);
     }
   }
+
+  private class InstrumentationContextMethodVisitor extends MethodVisitor {
+    // this data structure will remember last two LDC <class> instructions before
+    // InstrumentationContext.get() call
+    private final EvictingQueue<String> lastTwoClassConstants = EvictingQueue.create(2);
+
+    InstrumentationContextMethodVisitor(MethodVisitor methodVisitor) {
+      super(Opcodes.ASM7, methodVisitor);
+    }
+
+    @Override
+    public void visitInsn(int opcode) {
+      registerOpcode(opcode, null);
+      super.visitInsn(opcode);
+    }
+
+    @Override
+    public void visitIntInsn(int opcode, int operand) {
+      registerOpcode(opcode, null);
+      super.visitIntInsn(opcode, operand);
+    }
+
+    @Override
+    public void visitVarInsn(int opcode, int var) {
+      registerOpcode(opcode, null);
+      super.visitVarInsn(opcode, var);
+    }
+
+    @Override
+    public void visitTypeInsn(int opcode, String type) {
+      registerOpcode(opcode, null);
+      super.visitTypeInsn(opcode, type);
+    }
+
+    @Override
+    public void visitFieldInsn(int opcode, String owner, String name, String descriptor) {
+      registerOpcode(opcode, null);
+      super.visitFieldInsn(opcode, owner, name, descriptor);
+    }
+
+    @Override
+    public void visitMethodInsn(
+        int opcode, String owner, String name, String descriptor, boolean isInterface) {
+
+      Type methodType = Type.getMethodType(descriptor);
+      Type ownerType = Type.getType("L" + owner + ";");
+
+      // remember used context classes if this is an InstrumentationContext.get() call
+      if ("io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext"
+              .equals(ownerType.getClassName())
+          && "get".equals(name)
+          && methodType.getArgumentTypes().length == 2) {
+        // in case of invalid scenario (not using .class ref directly) don't store anything and
+        // clear the last LDC <class> stack
+        // note that FieldBackedProvider also check for an invalid context call in the runtime
+        if (lastTwoClassConstants.remainingCapacity() == 0) {
+          String className = lastTwoClassConstants.poll();
+          String contextClassName = lastTwoClassConstants.poll();
+          contextStoreClasses.put(className, contextClassName);
+        } else {
+          throw new MuzzleCompilationException(
+              "Invalid InstrumentationContext#get(Class, Class) usage: you cannot pass variables,"
+                  + " method parameters, compute classes; class references need to be passed"
+                  + " directly to the get() method");
+        }
+      }
+
+      registerOpcode(opcode, null);
+      super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+    }
+
+    @Override
+    public void visitJumpInsn(int opcode, Label label) {
+      registerOpcode(opcode, null);
+      super.visitJumpInsn(opcode, label);
+    }
+
+    @Override
+    public void visitLdcInsn(Object value) {
+      registerOpcode(Opcodes.LDC, value);
+      super.visitLdcInsn(value);
+    }
+
+    private void registerOpcode(int opcode, Object value) {
+      // check if this is an LDC <class> instruction; if so, remember the class that was used
+      // we need to remember last two LDC <class> instructions that were executed before
+      // InstrumentationContext.get() call
+      if (opcode == Opcodes.LDC) {
+        if (value instanceof Type) {
+          Type type = (Type) value;
+          if (type.getSort() == Type.OBJECT) {
+            lastTwoClassConstants.add(type.getClassName());
+            return;
+          }
+        }
+      }
+
+      // instruction other than LDC <class> visited; pop the first element if present - this will
+      // prevent adding wrong context key pairs in case of an invalid scenario
+      lastTwoClassConstants.poll();
+    }
+  }
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/ReferenceCollector.java

@@ -46,6 +46,7 @@ import net.bytebuddy.jar.asm.ClassReader;
 public class ReferenceCollector {
   private final Map<String, Reference> references = new LinkedHashMap<>();
   private final MutableGraph<String> helperSuperClassGraph = GraphBuilder.directed().build();
+  private final Map<String, String> contextStoreClasses = new LinkedHashMap<>();
   private final Set<String> visitedClasses = new HashSet<>();
   private final InstrumentationClassPredicate instrumentationClassPredicate;
 
@@ -142,6 +143,7 @@ public class ReferenceCollector {
         collectHelperClasses(
             isAdviceClass, visitedClassName, cv.getHelperClasses(), cv.getHelperSuperClasses());
 
+        contextStoreClasses.putAll(cv.getContextStoreClasses());
       } catch (IOException e) {
         throw new IllegalStateException("Error reading class " + visitedClassName, e);
       }
@@ -232,4 +234,8 @@ public class ReferenceCollector {
     }
     return helpersWithNoDeps;
   }
+
+  public Map<String, String> getContextStoreClasses() {
+    return contextStoreClasses;
+  }
 }

testing-common/integration-tests/src/main/java/context/ContextTestInstrumentationModule.java

@@ -47,15 +47,6 @@ public class ContextTestInstrumentationModule extends InstrumentationModule {
     return singletonList(new ContextTestInstrumentation());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    Map<String, String> store = new HashMap<>(3);
-    store.put("library.KeyClass", getClass().getName() + "$Context");
-    store.put("library.UntransformableKeyClass", getClass().getName() + "$Context");
-    store.put("library.DisabledKeyClass", getClass().getName() + "$Context");
-    return store;
-  }
-
   public static class ContextTestInstrumentation implements TypeInstrumentation {
     @Override
     public ElementMatcher<? super TypeDescription> typeMatcher() {
@@ -71,13 +62,6 @@ public class ContextTestInstrumentationModule extends InstrumentationModule {
       transformers.put(named("getContextCount"), GetApiUsageAdvice.class.getName());
       transformers.put(named("putContextCount"), PutApiUsageAdvice.class.getName());
       transformers.put(named("removeContextCount"), RemoveApiUsageAdvice.class.getName());
-      transformers.put(
-          named("incorrectKeyClassUsage"), IncorrectKeyClassContextApiUsageAdvice.class.getName());
-      transformers.put(
-          named("incorrectContextClassUsage"),
-          IncorrectContextClassContextApiUsageAdvice.class.getName());
-      transformers.put(
-          named("incorrectCallUsage"), IncorrectCallContextApiUsageAdvice.class.getName());
       return transformers;
     }
   }
@@ -142,30 +126,6 @@ public class ContextTestInstrumentationModule extends InstrumentationModule {
     }
   }
 
-  public static class IncorrectKeyClassContextApiUsageAdvice {
-    @Advice.OnMethodExit
-    public static void methodExit() {
-      InstrumentationContext.get(Object.class, Context.class);
-    }
-  }
-
-  public static class IncorrectContextClassContextApiUsageAdvice {
-    @Advice.OnMethodExit
-    public static void methodExit() {
-      InstrumentationContext.get(KeyClass.class, Object.class);
-    }
-  }
-
-  public static class IncorrectCallContextApiUsageAdvice {
-    @Advice.OnMethodExit
-    public static void methodExit() {
-      // Our instrumentation doesn't handle variables being passed to InstrumentationContext.get,
-      // so we make sure that this actually fails instrumentation.
-      Class clazz = null;
-      InstrumentationContext.get(clazz, Object.class);
-    }
-  }
-
   public static class Context {
     public static final ContextStore.Factory<Context> FACTORY = Context::new;
 

testing-common/integration-tests/src/test/groovy/context/FieldBackedProviderTest.groovy

@@ -15,9 +15,6 @@ import java.lang.reflect.Field
 import java.lang.reflect.Method
 import java.lang.reflect.Modifier
 import java.util.concurrent.atomic.AtomicReference
-import library.IncorrectCallUsageKeyClass
-import library.IncorrectContextClassUsageKeyClass
-import library.IncorrectKeyClassUsageKeyClass
 import library.KeyClass
 import library.UntransformableKeyClass
 import net.bytebuddy.agent.ByteBuddyAgent
@@ -187,21 +184,6 @@ class FieldBackedProviderTest extends AgentInstrumentationSpecification {
     new KeyClass().incrementContextCount() == 1
     new UntransformableKeyClass().incrementContextCount() == 1
   }
-
-  def "incorrect key class usage fails at class load time"() {
-    expect:
-    !new IncorrectKeyClassUsageKeyClass().isInstrumented()
-  }
-
-  def "incorrect context class usage fails at class load time"() {
-    expect:
-    !new IncorrectContextClassUsageKeyClass().isInstrumented()
-  }
-
-  def "incorrect call usage fails at class load time"() {
-    expect:
-    !new IncorrectCallUsageKeyClass().isInstrumented()
-  }
 }
 
 

testing-common/library-for-integration-tests/src/main/java/library/IncorrectCallUsageKeyClass.java

@@ -1,17 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package library;
-
-public class IncorrectCallUsageKeyClass {
-  public boolean isInstrumented() {
-    return false;
-  }
-
-  public int incorrectCallUsage() {
-    // instrumentation will not apply to this class because advice incorrectly uses context api
-    return -1;
-  }
-}

testing-common/library-for-integration-tests/src/main/java/library/IncorrectContextClassUsageKeyClass.java

@@ -1,17 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package library;
-
-public class IncorrectContextClassUsageKeyClass {
-  public boolean isInstrumented() {
-    return false;
-  }
-
-  public int incorrectContextClassUsage() {
-    // instrumentation will not apply to this class because advice incorrectly uses context api
-    return -1;
-  }
-}

testing-common/library-for-integration-tests/src/main/java/library/IncorrectKeyClassUsageKeyClass.java

@@ -1,17 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package library;
-
-public class IncorrectKeyClassUsageKeyClass {
-  public boolean isInstrumented() {
-    return false;
-  }
-
-  public int incorrectKeyClassUsage() {
-    // instrumentation will not apply to this class because advice incorrectly uses context api
-    return -1;
-  }
-}

testing-common/src/test/groovy/muzzle/ReferenceCollectorTest.groovy

@@ -14,9 +14,12 @@ import static muzzle.TestClasses.LdcAdvice
 import static muzzle.TestClasses.MethodBodyAdvice
 
 import external.instrumentation.ExternalHelper
+import io.opentelemetry.context.Context
+import io.opentelemetry.instrumentation.InstrumentationContextTestClasses
 import io.opentelemetry.instrumentation.OtherTestHelperClasses
 import io.opentelemetry.instrumentation.TestHelperClasses
 import io.opentelemetry.javaagent.tooling.muzzle.Reference
+import io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCompilationException
 import io.opentelemetry.javaagent.tooling.muzzle.collector.ReferenceCollector
 import spock.lang.Specification
 import spock.lang.Unroll
@@ -255,6 +258,33 @@ class ReferenceCollectorTest extends Specification {
     collector.sortedHelperClasses.isEmpty()
   }
 
+  def "should collect context store classes"() {
+    when:
+    def collector = new ReferenceCollector({ false })
+    collector.collectReferencesFromAdvice(InstrumentationContextTestClasses.ValidAdvice.name)
+
+    then:
+    def contextStore = collector.getContextStoreClasses()
+    contextStore == [
+      (InstrumentationContextTestClasses.Key1.name): Context.name,
+      (InstrumentationContextTestClasses.Key2.name): Context.name
+    ]
+  }
+
+  def "should not collect context store classes for invalid scenario: #desc"() {
+    when:
+    def collector = new ReferenceCollector({ false })
+    collector.collectReferencesFromAdvice(adviceClassName)
+
+    then:
+    thrown(MuzzleCompilationException)
+
+    where:
+    desc                                                                        | adviceClassName
+    "passing arbitrary variables or parameters to InstrumentationContext.get()" | InstrumentationContextTestClasses.NotUsingClassRefAdvice.name
+    "storing class ref in a local var"                                          | InstrumentationContextTestClasses.PassingVariableAdvice.name
+  }
+
   private static assertHelperSuperClassMethod(Reference reference, boolean isAbstract) {
     assertMethod reference, 'abstractMethod', '()I',
       VisibilityFlag.PROTECTED,

testing-common/src/test/java/io/opentelemetry/instrumentation/InstrumentationContextTestClasses.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+
+public class InstrumentationContextTestClasses {
+  public static class ValidAdvice {
+    public static void advice() {
+      Runnable.class.getName();
+      InstrumentationContext.get(Key1.class, Context.class);
+      Key2.class.getName();
+      Key1.class.getName();
+      InstrumentationContext.get(Key2.class, Context.class);
+    }
+  }
+
+  public static class NotUsingClassRefAdvice {
+    public static void advice(Class<?> key, Class<?> context) {
+      Key2.class.getName();
+      Key1.class.getName();
+      InstrumentationContext.get(key, context);
+    }
+  }
+
+  public static class PassingVariableAdvice {
+    public static void advice() {
+      Class<?> context = Context.class;
+      InstrumentationContext.get(Key1.class, context);
+    }
+  }
+
+  public static class Key1 {}
+
+  public static class Key2 {}
+}