KAFKA-14654: Connector classes should statically initialize with plugin classloader

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoader.java

@@ -400,7 +400,7 @@ private <T> Collection<PluginDesc<T>> getPluginDesc(
         Collection<PluginDesc<T>> result = new ArrayList<>();
         for (Class<? extends T> plugin : plugins) {
             if (PluginUtils.isConcrete(plugin)) {
-                try {
+                try (LoaderSwap loaderSwap = withClassLoader(loader)) {
                     result.add(pluginDesc(plugin, versionFor(plugin), loader));
                 } catch (ReflectiveOperationException | LinkageError e) {
                     log.error("Failed to discover {}: Unable to instantiate {}{}", klass.getSimpleName(), plugin.getSimpleName(), reflectiveErrorDescription(e), e);
@@ -419,11 +419,10 @@ private <T> PluginDesc<T> pluginDesc(Class<? extends T> plugin, String version,
 
     @SuppressWarnings("unchecked")
     private <T> Collection<PluginDesc<T>> getServiceLoaderPluginDesc(Class<T> klass, ClassLoader loader) {
-        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
         Collection<PluginDesc<T>> result = new ArrayList<>();
-        try {
-            ServiceLoader<T> serviceLoader = ServiceLoader.load(klass, loader);
-            for (Iterator<T> iterator = serviceLoader.iterator(); iterator.hasNext(); ) {
+        ServiceLoader<T> serviceLoader = ServiceLoader.load(klass, loader);
+        for (Iterator<T> iterator = serviceLoader.iterator(); iterator.hasNext(); ) {
+            try (LoaderSwap loaderSwap = withClassLoader(loader)) {
                 T pluginImpl;
                 try {
                     pluginImpl = iterator.next();
@@ -434,8 +433,6 @@ private <T> Collection<PluginDesc<T>> getServiceLoaderPluginDesc(Class<T> klass,
                 result.add(pluginDesc((Class<? extends T>) pluginImpl.getClass(),
                     versionFor(pluginImpl), loader));
             }
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
         }
         return result;
     }
@@ -473,6 +470,16 @@ private static String reflectiveErrorDescription(Throwable t) {
         }
     }
 
+    public LoaderSwap withClassLoader(ClassLoader loader) {
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
+        try {
+            return new LoaderSwap(savedLoader);
+        } catch (Throwable t) {
+            Plugins.compareAndSwapLoaders(savedLoader);
+            throw t;
+        }
+    }
+
     @Override
     protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
         String fullName = aliases.getOrDefault(name, name);

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java

@@ -30,6 +30,7 @@
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.ConnectException;
@@ -295,7 +296,7 @@ public void newPluginShouldServiceLoadWithPluginClassLoader() {
         // Assert that the service loaded subclass is found in both environments
         assertTrue(samples.containsKey("ServiceLoadedSubclass.static"));
         assertTrue(samples.containsKey("ServiceLoadedSubclass.dynamic"));
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -306,9 +307,7 @@ public void newPluginShouldInstantiateWithPluginClassLoader() {
             Converter.class
         );
 
-        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
-        Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -334,7 +333,7 @@ public void newConverterShouldConfigureWithPluginClassLoader() {
         assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
         Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
         assertTrue(samples.containsKey("configure"));
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -357,7 +356,7 @@ public void newConfigProviderShouldConfigureWithPluginClassLoader() {
         assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
         Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
         assertTrue(samples.containsKey("configure"));
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -377,7 +376,17 @@ public void newHeaderConverterShouldConfigureWithPluginClassLoader() {
         assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
         Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
         assertTrue(samples.containsKey("configure")); // HeaderConverter::configure was called
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
+    }
+
+    @Test
+    public void newConnectorShouldInstantiateWithPluginClassLoader() {
+        Connector plugin = plugins.newConnector(TestPlugin.SAMPLING_CONNECTOR.className());
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
+        assertTrue(samples.containsKey("<init>")); // constructor was called
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -393,7 +402,7 @@ public void newPluginsShouldConfigureWithPluginClassLoader() {
         assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
         Map<String, SamplingTestPlugin> samples = ((SamplingTestPlugin) plugin).flatten();
         assertTrue(samples.containsKey("configure")); // Configurable::configure was called
-        assertPluginClassLoaderAlwaysActive(samples);
+        assertPluginClassLoaderAlwaysActive(plugin);
     }
 
     @Test
@@ -460,19 +469,23 @@ private void assertClassLoaderReadsVersionFromResource(
                 converter.toConnectData(null, null).value());
     }
 
-    public static void assertPluginClassLoaderAlwaysActive(Map<String, SamplingTestPlugin> samples) {
-        for (Entry<String, SamplingTestPlugin> e : samples.entrySet()) {
-            String sampleName = "\"" + e.getKey() + "\" (" + e.getValue() + ")";
-            assertInstanceOf(
-                PluginClassLoader.class,
-                e.getValue().staticClassloader(),
-                sampleName + " has incorrect static classloader"
-            );
-            assertInstanceOf(
-                PluginClassLoader.class,
-                e.getValue().classloader(),
-                sampleName + " has incorrect dynamic classloader"
-            );
+    public static void assertPluginClassLoaderAlwaysActive(Object plugin) {
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        for (SamplingTestPlugin instance : ((SamplingTestPlugin) plugin).allInstances()) {
+            Map<String, SamplingTestPlugin> samples = instance.flatten();
+            for (Entry<String, SamplingTestPlugin> e : samples.entrySet()) {
+                String sampleName = "\"" + e.getKey() + "\" (" + e.getValue() + ")";
+                assertInstanceOf(
+                        PluginClassLoader.class,
+                        e.getValue().staticClassloader(),
+                        sampleName + " has incorrect static classloader"
+                );
+                assertInstanceOf(
+                        PluginClassLoader.class,
+                        e.getValue().classloader(),
+                        sampleName + " has incorrect dynamic classloader"
+                );
+            }
         }
     }
 

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/SamplingTestPlugin.java

@@ -19,36 +19,44 @@
 
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 
 /**
  * Base class for plugins so we can sample information about their initialization
  */
-public abstract class SamplingTestPlugin {
+public interface SamplingTestPlugin {
 
     /**
      * @return the ClassLoader used to statically initialize this plugin class
      */
-    public abstract ClassLoader staticClassloader();
+    ClassLoader staticClassloader();
 
     /**
      * @return the ClassLoader used to initialize this plugin instance
      */
-    public abstract ClassLoader classloader();
+    ClassLoader classloader();
+
+    /**
+     * @return All known instances of this class, including this instance.
+     */
+    default List<SamplingTestPlugin> allInstances() {
+        return Collections.singletonList(this);
+    }
 
     /**
      * @return a group of other SamplingTestPlugin instances known by this plugin
      * This should only return direct children, and not reference this instance directly
      */
-    public Map<String, SamplingTestPlugin> otherSamples() {
+    default Map<String, SamplingTestPlugin> otherSamples() {
         return Collections.emptyMap();
     }
 
     /**
      * @return a flattened list of child samples including this entry keyed as "this"
      */
-    public Map<String, SamplingTestPlugin> flatten() {
+    default Map<String, SamplingTestPlugin> flatten() {
         Map<String, SamplingTestPlugin> out = new HashMap<>();
         Map<String, SamplingTestPlugin> otherSamples = otherSamples();
         if (otherSamples != null) {
@@ -71,7 +79,7 @@ public Map<String, SamplingTestPlugin> flatten() {
      * Stores only the last invocation of each method if there are multiple invocations.
      * @param samples The collection of samples to which this method call should be added
      */
-    public void logMethodCall(Map<String, SamplingTestPlugin> samples) {
+    default void logMethodCall(Map<String, SamplingTestPlugin> samples) {
         StackTraceElement[] stackTraces = Thread.currentThread().getStackTrace();
         if (stackTraces.length < 2) {
             return;
@@ -88,7 +96,7 @@ public void logMethodCall(Map<String, SamplingTestPlugin> samples) {
         ));
     }
 
-    public static class MethodCallSample extends SamplingTestPlugin {
+    class MethodCallSample implements SamplingTestPlugin {
 
         private final StackTraceElement caller;
         private final ClassLoader staticClassLoader;

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/TestPlugins.java

@@ -94,6 +94,11 @@ public enum TestPlugin {
          * which samples information about its method calls.
          */
         SAMPLING_CONFIG_PROVIDER("sampling-config-provider", "test.plugins.SamplingConfigProvider"),
+        /**
+         * A {@link org.apache.kafka.connect.sink.SinkConnector}
+         * which samples information about its method calls.
+         */
+        SAMPLING_CONNECTOR("sampling-connector", "test.plugins.SamplingConnector"),
         /**
          * A plugin which uses a {@link java.util.ServiceLoader}
          * to load internal classes, and samples information about their initialization.

connect/runtime/src/test/resources/test-plugins/aliased-static-field/test/plugins/AliasedStaticField.java

@@ -30,7 +30,7 @@
  * <p>Samples data about its initialization environment for later analysis.
  * Samples are shared between instances of the same class in a static variable.
  */
-public class AliasedStaticField extends SamplingTestPlugin implements Converter {
+public class AliasedStaticField implements SamplingTestPlugin, Converter {
 
   private static final Map<String, SamplingTestPlugin> SAMPLES;
   private static final ClassLoader STATIC_CLASS_LOADER;

connect/runtime/src/test/resources/test-plugins/always-throw-exception/test/plugins/AlwaysThrowException.java

@@ -20,7 +20,6 @@
 import java.util.Map;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
-import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
 import org.apache.kafka.connect.storage.Converter;
 
 /**

connect/runtime/src/test/resources/test-plugins/multiple-plugins-in-jar/test/plugins/ThingOne.java

@@ -21,7 +21,6 @@
 import java.util.Map;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
-import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
 import org.apache.kafka.connect.storage.Converter;
 
 /**

connect/runtime/src/test/resources/test-plugins/multiple-plugins-in-jar/test/plugins/ThingTwo.java

@@ -21,7 +21,6 @@
 import java.util.Map;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
-import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
 import org.apache.kafka.connect.runtime.isolation.TestPlugins;
 import org.apache.kafka.connect.storage.Converter;
 

connect/runtime/src/test/resources/test-plugins/sampling-config-provider/test/plugins/SamplingConfigProvider.java

@@ -17,9 +17,13 @@
 
 package test.plugins;
 
+import java.util.Collections;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Set;
 import java.util.Map;
 import java.util.HashMap;
+
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.common.config.ConfigData;
 import org.apache.kafka.common.config.ConfigChangeCallback;
@@ -34,14 +38,16 @@
  * See {@link org.apache.kafka.connect.runtime.isolation.TestPlugins}.
  * <p>Samples data about its initialization environment for later analysis.
  */
-public class SamplingConfigProvider extends SamplingTestPlugin implements ConfigProvider {
+public final class SamplingConfigProvider implements SamplingTestPlugin, ConfigProvider {
 
   private static final ClassLoader STATIC_CLASS_LOADER;
+  private static List<SamplingTestPlugin> instances;
   private final ClassLoader classloader;
   private Map<String, SamplingTestPlugin> samples;
 
   static {
     STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+    instances = Collections.synchronizedList(new ArrayList<>());
   }
 
   {
@@ -61,6 +67,11 @@ public ConfigData get(String path, Set<String> keys) {
     return null;
   }
 
+  public SamplingConfigProvider() {
+    logMethodCall(samples);
+    instances.add(this);
+  }
+
   @Override
   public void subscribe(String path, Set<String> keys, ConfigChangeCallback callback) {
     logMethodCall(samples);
@@ -100,4 +111,10 @@ public ClassLoader classloader() {
   public Map<String, SamplingTestPlugin> otherSamples() {
     return samples;
   }
+
+
+  @Override
+  public List<SamplingTestPlugin> allInstances() {
+    return instances;
+  }
 }

connect/runtime/src/test/resources/test-plugins/sampling-configurable/test/plugins/SamplingConfigurable.java

@@ -17,8 +17,12 @@
 
 package test.plugins;
 
+import java.util.Collections;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.HashMap;
+
 import org.apache.kafka.common.Configurable;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
@@ -30,21 +34,28 @@
  * See {@link org.apache.kafka.connect.runtime.isolation.TestPlugins}.
  * <p>Samples data about its initialization environment for later analysis.
  */
-public class SamplingConfigurable extends SamplingTestPlugin implements Converter, Configurable {
+public final class SamplingConfigurable implements SamplingTestPlugin, Converter, Configurable {
 
   private static final ClassLoader STATIC_CLASS_LOADER;
+  private static List<SamplingTestPlugin> instances;
   private final ClassLoader classloader;
   private Map<String, SamplingTestPlugin> samples;
 
   static {
     STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+    instances = Collections.synchronizedList(new ArrayList<>());
   }
 
   {
     samples = new HashMap<>();
     classloader = Thread.currentThread().getContextClassLoader();
   }
 
+  public SamplingConfigurable() {
+    logMethodCall(samples);
+    instances.add(this);
+  }
+
   @Override
   public void configure(final Map<String, ?> configs) {
     logMethodCall(samples);
@@ -78,4 +89,9 @@ public ClassLoader classloader() {
   public Map<String, SamplingTestPlugin> otherSamples() {
     return samples;
   }
+
+  @Override
+  public List<SamplingTestPlugin> allInstances() {
+    return instances;
+  }
 }