From b4326d91ac23609968fbc32eb65abba94f331f1f Mon Sep 17 00:00:00 2001
From: Pavol Loffay <p.loffay@gmail.com>
Date: Fri, 13 Nov 2020 20:21:40 +0100
Subject: [PATCH] Add Bytebuddy agent builder customizer SPI (#1613)

---
 javaagent-spi/javaagent-spi.gradle            |  1 +
 .../spi/ByteBuddyAgentCustomizer.java         | 30 +++++++++++++++++++
 .../javaagent/tooling/AgentInstaller.java     | 17 +++++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 javaagent-spi/src/main/java/io/opentelemetry/javaagent/spi/ByteBuddyAgentCustomizer.java

diff --git a/javaagent-spi/javaagent-spi.gradle b/javaagent-spi/javaagent-spi.gradle
index 1944baca5e..b9af4093f7 100644
--- a/javaagent-spi/javaagent-spi.gradle
+++ b/javaagent-spi/javaagent-spi.gradle
@@ -5,4 +5,5 @@ apply from: "$rootDir/gradle/publish.gradle"
 
 dependencies {
   compileOnly deps.opentelemetrySdk
+  compileOnly deps.bytebuddy
 }
\ No newline at end of file
diff --git a/javaagent-spi/src/main/java/io/opentelemetry/javaagent/spi/ByteBuddyAgentCustomizer.java b/javaagent-spi/src/main/java/io/opentelemetry/javaagent/spi/ByteBuddyAgentCustomizer.java
new file mode 100644
index 0000000000..94f423d049
--- /dev/null
+++ b/javaagent-spi/src/main/java/io/opentelemetry/javaagent/spi/ByteBuddyAgentCustomizer.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.spi;
+
+import java.lang.instrument.Instrumentation;
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+/**
+ * {@link ByteBuddyAgentCustomizer} customizes ByteBuddy agent builder right before the agent is
+ * installed - {@link AgentBuilder#installOn(Instrumentation)}. This SPI can be used to customize
+ * {@link AgentBuilder} for vendor specific needs. For example install custom listeners or exclude
+ * classes. Use this SPI carefully because it can change {@link net.bytebuddy.ByteBuddy} behaviour.
+ *
+ * <p>This is a service provider interface that requires implementations to be registered in {@code
+ * META-INF/services} folder.
+ */
+public interface ByteBuddyAgentCustomizer {
+
+  /**
+   * Customize ByteBuddy agent builder before {@link AgentBuilder#installOn(Instrumentation)} is
+   * called.
+   *
+   * @param agentBuilder ByteBuddy agent from {@code AgentInstaller}.
+   * @return customized agent builder.
+   */
+  AgentBuilder customize(AgentBuilder agentBuilder);
+}
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 38a8a0f2fb..49c2ba2f9a 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
@@ -18,6 +18,7 @@ import io.opentelemetry.instrumentation.api.internal.BootstrapPackagePrefixesHol
 import io.opentelemetry.javaagent.instrumentation.api.OpenTelemetrySdkAccess;
 import io.opentelemetry.javaagent.instrumentation.api.SafeServiceLoader;
 import io.opentelemetry.javaagent.spi.BootstrapPackagesProvider;
+import io.opentelemetry.javaagent.spi.ByteBuddyAgentCustomizer;
 import io.opentelemetry.javaagent.tooling.config.ConfigInitializer;
 import io.opentelemetry.javaagent.tooling.context.FieldBackedProvider;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -29,6 +30,7 @@ import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.stream.Collectors;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
@@ -156,10 +158,25 @@ public class AgentInstaller {
       }
     }
 
+    agentBuilder = customizeByteBuddyAgent(agentBuilder);
     log.debug("Installed {} instrumenter(s)", numInstrumenters);
     return agentBuilder.installOn(inst);
   }
 
+  private static AgentBuilder customizeByteBuddyAgent(AgentBuilder agentBuilder) {
+    Iterable<ByteBuddyAgentCustomizer> agentCustomizers = loadByteBuddyAgentCustomizers();
+    for (ByteBuddyAgentCustomizer agentCustomizer : agentCustomizers) {
+      log.debug("Applying agent builder customizer {}", agentCustomizer.getClass().getName());
+      agentBuilder = agentCustomizer.customize(agentBuilder);
+    }
+    return agentBuilder;
+  }
+
+  private static Iterable<ByteBuddyAgentCustomizer> loadByteBuddyAgentCustomizers() {
+    return ServiceLoader.load(
+        ByteBuddyAgentCustomizer.class, AgentInstaller.class.getClassLoader());
+  }
+
   private static List<InstrumentationModule> loadInstrumentationModules() {
     return SafeServiceLoader.load(
             InstrumentationModule.class, AgentInstaller.class.getClassLoader())