From 2311d71def5e86372d76596273fdcfe1f0ba0e1f Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Thu, 7 Apr 2022 00:08:16 +0300
Subject: [PATCH] add an spi to get access to bootstrap proxy from muzzle
 module

---
 .../javaagent/tooling/AgentInstaller.java     |  1 -
 .../bootstrap/BootstrapProxyProviderImpl.java | 19 +++++++++++++++
 .../tooling/muzzle/AgentTooling.java          | 23 +++++++++++--------
 .../muzzle/BootstrapProxyProvider.java        | 13 +++++++++++
 4 files changed, 46 insertions(+), 10 deletions(-)
 create mode 100644 javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapProxyProviderImpl.java
 create mode 100644 muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/BootstrapProxyProvider.java

diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
index 29ad0189db4b..4e1108a1c42d 100644
--- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
+++ b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
@@ -137,7 +137,6 @@ public static ResettableClassFileTransformer installBytebuddyAgent(
       runBeforeAgentListeners(agentListeners, config, autoConfiguredSdk);
     }
 
-    AgentTooling.init(Utils.getBootstrapProxy());
     AgentBuilder agentBuilder =
         new AgentBuilder.Default()
             .disableClassFormatChanges()
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapProxyProviderImpl.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapProxyProviderImpl.java
new file mode 100644
index 000000000000..d4b301a5cf42
--- /dev/null
+++ b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapProxyProviderImpl.java
@@ -0,0 +1,19 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.bootstrap;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.tooling.Utils;
+import io.opentelemetry.javaagent.tooling.muzzle.BootstrapProxyProvider;
+
+@AutoService(BootstrapProxyProvider.class)
+public class BootstrapProxyProviderImpl implements BootstrapProxyProvider {
+
+  @Override
+  public ClassLoader getBootstrapProxy() {
+    return Utils.getBootstrapProxy();
+  }
+}
diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java
index 7f69c138368e..a2dd7f8e749b 100644
--- a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java
+++ b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java
@@ -5,6 +5,8 @@
 
 package io.opentelemetry.javaagent.tooling.muzzle;
 
+import java.util.Iterator;
+import java.util.ServiceLoader;
 import net.bytebuddy.agent.builder.AgentBuilder;
 
 /**
@@ -14,7 +16,8 @@
  */
 public final class AgentTooling {
 
-  private static volatile ClassLoader bootstrapProxy;
+  private static final AgentBuilder.PoolStrategy POOL_STRATEGY =
+      new AgentCachingPoolStrategy(locationStrategy(getBootstrapProxy()));
 
   public static AgentLocationStrategy locationStrategy() {
     return locationStrategy(null);
@@ -25,17 +28,19 @@ public static AgentLocationStrategy locationStrategy(ClassLoader bootstrapProxy)
   }
 
   public static AgentBuilder.PoolStrategy poolStrategy() {
-    return PoolStrategyHolder.POOL_STRATEGY;
+    return POOL_STRATEGY;
   }
 
-  public static void init(ClassLoader classLoader) {
-    bootstrapProxy = classLoader;
+  private static ClassLoader getBootstrapProxy() {
+    Iterator<BootstrapProxyProvider> iterator =
+        ServiceLoader.load(BootstrapProxyProvider.class).iterator();
+    if (iterator.hasNext()) {
+      BootstrapProxyProvider bootstrapProxyProvider = iterator.next();
+      return bootstrapProxyProvider.getBootstrapProxy();
+    }
+
+    return null;
   }
 
   private AgentTooling() {}
-
-  private static class PoolStrategyHolder {
-    private static final AgentBuilder.PoolStrategy POOL_STRATEGY =
-        new AgentCachingPoolStrategy(locationStrategy(bootstrapProxy));
-  }
 }
diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/BootstrapProxyProvider.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/BootstrapProxyProvider.java
new file mode 100644
index 000000000000..d3e1548dba1e
--- /dev/null
+++ b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/BootstrapProxyProvider.java
@@ -0,0 +1,13 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.muzzle;
+
+/** This SPI can be used to find a class loader which can be used to look up bootstrap resources. */
+public interface BootstrapProxyProvider {
+
+  /** Provide a class loader which can be used to look up bootstrap resources. */
+  ClassLoader getBootstrapProxy();
+}