Update ShadowJar Plugin

build.gradle

@@ -5,13 +5,15 @@
 
 buildscript {
   repositories {
+    gradlePluginPortal()
     jcenter()
     maven {
       url "https://plugins.gradle.org/m2/"
     }
   }
+
   dependencies {
-    classpath 'com.github.jengelman.gradle.plugins:shadow:2.0.4'
+    classpath 'com.github.johnrengelman.shadow:com.github.johnrengelman.shadow.gradle.plugin:7.1.2'
     classpath 'com.github.alisiikh:gradle-scalastyle-plugin:3.4.1'
     classpath "io.github.gradle-nexus:publish-plugin:1.0.0"
     classpath "org.shipkit:shipkit-auto-version:1.1.1"

transportable-udfs-examples/build.gradle

@@ -25,8 +25,9 @@ allprojects {
 subprojects {
   buildscript {
     repositories {
-      mavenCentral()
+      gradlePluginPortal()
       jcenter()
+      mavenCentral()
     }
   }
   repositories {

transportable-udfs-plugin/build.gradle

@@ -4,12 +4,16 @@ plugins {
   id 'signing'
 }
 
+repositories { 
+  gradlePluginPortal()
+}
+
 dependencies {
   implementation project(':transportable-udfs-api')
   implementation project(':transportable-udfs-codegen')
   implementation ('com.google.guava:guava:24.1-jre')
   implementation ('com.google.code.gson:gson:2.8.5')
-  implementation ('com.github.jengelman.gradle.plugins:shadow:5.2.0')
+  implementation ('com.github.johnrengelman.shadow:com.github.johnrengelman.shadow.gradle.plugin:7.1.2')
   testImplementation('org.spockframework:spock-core:1.1-groovy-2.4') {
     exclude group: 'org.codehaus.groovy'
   }

transportable-udfs-plugin/src/main/java/com/linkedin/transport/plugin/TransportPlugin.java

@@ -57,7 +57,7 @@ public void apply(Project project) {
     project.getPlugins().withType(JavaPlugin.class, (javaPlugin) -> {
       project.getPlugins().apply(ScalaPlugin.class);
       project.getPlugins().apply(DistributionPlugin.class);
-      project.getConfigurations().create(ShadowBasePlugin.getCONFIGURATION_NAME());
+      project.getConfigurations().create(ShadowBasePlugin.CONFIGURATION_NAME);
 
       JavaPluginConvention javaConvention = project.getConvention().getPlugin(JavaPluginConvention.class);
       SourceSet mainSourceSet = javaConvention.getSourceSets().getByName(extension.mainSourceSetName);

transportable-udfs-plugin/src/main/java/com/linkedin/transport/plugin/packaging/DistributionPackaging.java

@@ -57,9 +57,9 @@ public List<TaskProvider<? extends Task>> configurePackagingTasks(Project projec
     // Explicitly set classifiers for the created distributions or else leads to Maven packaging issues due to multiple
     // artifacts with the same classifier
     project.getTasks().named(platform.getName() + "DistTar", Tar.class, tar -> tar.setClassifier(platform.getName()));
-    project.getArtifacts().add(ShadowBasePlugin.getCONFIGURATION_NAME(), project.getTasks().named(platform.getName() + "DistTar", Tar.class));
+    project.getArtifacts().add(ShadowBasePlugin.CONFIGURATION_NAME, project.getTasks().named(platform.getName() + "DistTar", Tar.class));
     project.getTasks().named(platform.getName() + "DistZip", Zip.class, zip -> zip.setClassifier(platform.getName()));
-    project.getArtifacts().add(ShadowBasePlugin.getCONFIGURATION_NAME(), project.getTasks().named(platform.getName() + "DistZip", Zip.class));
+    project.getArtifacts().add(ShadowBasePlugin.CONFIGURATION_NAME, project.getTasks().named(platform.getName() + "DistZip", Zip.class));
     return ImmutableList.of(project.getTasks().named(platform.getName() + "DistTar", Tar.class),
         project.getTasks().named(platform.getName() + "DistZip", Zip.class));
   }

transportable-udfs-plugin/src/main/java/com/linkedin/transport/plugin/packaging/ShadedJarPackaging.java

@@ -7,36 +7,55 @@
 
 import com.github.jengelman.gradle.plugins.shadow.ShadowBasePlugin;
 import com.github.jengelman.gradle.plugins.shadow.ShadowJavaPlugin;
+import com.github.jengelman.gradle.plugins.shadow.relocation.SimpleRelocator;
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableMap;
 import com.linkedin.transport.plugin.Platform;
-import com.linkedin.transport.plugin.tasks.ShadeTask;
-import java.util.List;
 import org.gradle.api.Project;
 import org.gradle.api.Task;
+import org.gradle.api.artifacts.Configuration;
 import org.gradle.api.component.AdhocComponentWithVariants;
+import org.gradle.api.component.ConfigurationVariantDetails;
+import org.gradle.api.file.FileCollection;
+import org.gradle.api.logging.Logger;
+import org.gradle.api.logging.Logging;
 import org.gradle.api.tasks.SourceSet;
 import org.gradle.api.tasks.TaskProvider;
 import org.gradle.api.tasks.bundling.Jar;
 
-import static com.linkedin.transport.plugin.ConfigurationType.*;
-import static com.linkedin.transport.plugin.SourceSetUtils.*;
+import java.io.File;
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.jar.JarFile;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.zip.ZipEntry;
+
+import static com.linkedin.transport.plugin.ConfigurationType.RUNTIME_CLASSPATH;
+import static com.linkedin.transport.plugin.SourceSetUtils.getConfigurationForSourceSet;
 
 
 /**
  * A {@link Packaging} class which generates Shaded JARs containing all runtime dependencies using the
- * {@link ShadeTask}
+ * {@link ShadowJar}
  */
 public class ShadedJarPackaging implements Packaging {
 
+  private static final Logger LOG = Logging.getLogger(ShadedJarPackaging.class);
+
   private final List<String> _excludedDependencies;
   private final List<String> _classesToNotShade;
 
   /**
    * Creates a {@link ShadedJarPackaging} object with the given configuration
    *
-   * @param excludedDependencies Dependencies to be excluded form the shaded JAR. See {@link ShadeTask#getExcludedDependencies()}
-   * @param classesToNotShade Classes to be kept unshaded in the shaded JAR. See {@link ShadeTask#getDoNotShade()}
+   * @param excludedDependencies Dependencies to be excluded form the shaded JAR.
+   * @param classesToNotShade Classes to be kept unshaded in the shaded JAR.
    */
   public ShadedJarPackaging(List<String> excludedDependencies, List<String> classesToNotShade) {
     _excludedDependencies = excludedDependencies;
@@ -46,43 +65,98 @@ public ShadedJarPackaging(List<String> excludedDependencies, List<String> classe
   @Override
   public List<TaskProvider<? extends Task>> configurePackagingTasks(Project project, Platform platform,
       SourceSet platformSourceSet, SourceSet mainSourceSet) {
-    TaskProvider<ShadeTask> shadeTask = createShadeTask(project, platform, platformSourceSet, mainSourceSet);
-    shadeTask.configure(task -> {
-      if (_excludedDependencies != null) {
-        task.setExcludedDependencies(ImmutableSet.copyOf(_excludedDependencies));
-      }
-      if (_classesToNotShade != null) {
-        task.setDoNotShade(_classesToNotShade);
-      }
-    });
-    return ImmutableList.of(shadeTask);
+    return ImmutableList.of(createShadeTask(project, platform, platformSourceSet, mainSourceSet));
   }
 
   /**
-   * Creates a {@link ShadeTask} which generates a shaded JAR containing all runtime dependencies of the platform's
-   * {@link SourceSet}
-   *
-   * TODO: This code is borrowed from the Shade plugin. Call the functionality residing in the Shade plugin once it is
-   * available
+   * Creates a {@link ShadowJar} which generates a shaded JAR containing all runtime dependencies of the platform's
+   * {@link SourceSet}.
    */
-  private TaskProvider<ShadeTask> createShadeTask(Project project, Platform platform, SourceSet sourceSet,
+  private TaskProvider<ShadowJar> createShadeTask(Project project, Platform platform, SourceSet sourceSet,
       SourceSet mainSourceSet) {
-    TaskProvider<ShadeTask> shadeTask =
-        project.getTasks().register(sourceSet.getTaskName("shade", "Jar"), ShadeTask.class, task -> {
-          task.setGroup(ShadowJavaPlugin.getSHADOW_GROUP());
-          task.setDescription("Create a combined JAR of " + platform.getName() + " output and runtime dependencies");
+    TaskProvider<ShadowJar> shadeTask =
+        project.getTasks().register(sourceSet.getTaskName("shade", "Jar"), ShadowJar.class, task -> {
+          task.setGroup(ShadowJavaPlugin.SHADOW_GROUP);
+          task.setDescription("Create a combined JAR of " + platform.getName() + " output and runtime dependencies.");
           task.setClassifier(platform.getName());
           task.getManifest()
               .inheritFrom(project.getTasks().named(mainSourceSet.getJarTaskName(), Jar.class).get().getManifest());
+
           task.from(sourceSet.getOutput());
-          task.setConfigurations(ImmutableList.of(getConfigurationForSourceSet(project, sourceSet, RUNTIME_CLASSPATH)));
+          task.setConfigurations(ImmutableList.of(excludeDependencies(getConfigurationForSourceSet(project, sourceSet, RUNTIME_CLASSPATH))));
           task.exclude("META-INF/INDEX.LIST", "META-INF/*.SF", "META-INF/*.DSA", "META-INF/*.RSA");
-        });
 
-    String configuration = ShadowBasePlugin.getCONFIGURATION_NAME();
-    project.getArtifacts().add(configuration, shadeTask);
+          task.relocate(new SimpleRelocator(null, shadePrefix(task.getProject()) + "/", ImmutableList.of("*.class"), _classesToNotShade) {
+            private final Set<String> classPathsToShade = task.getConfigurations().stream()
+                .flatMap(files -> classesInConfiguration(files).stream())
+                .collect(Collectors.toSet());
+            private final Set<String> classNamesToShade = classPathsToShade.stream()
+                .map(path -> path.replace('/', '.'))
+                .collect(Collectors.toSet());
+
+            @Override
+            public boolean canRelocatePath(String path) {
+              return classPathsToShade.contains(path) && super.canRelocatePath(path);
+            }
+
+            @Override
+            public boolean canRelocateClass(String className) {
+              return classNamesToShade.contains(className) && super.canRelocateClass(className);
+            }
+          });
+        });
+
+    project.getArtifacts().add(ShadowBasePlugin.CONFIGURATION_NAME, shadeTask);
     AdhocComponentWithVariants java = project.getComponents().withType(AdhocComponentWithVariants.class).getByName("java");
-    java.addVariantsFromConfiguration(project.getConfigurations().getByName(configuration), v -> v.mapToOptional());
+    java.addVariantsFromConfiguration(project.getConfigurations().getByName(ShadowBasePlugin.CONFIGURATION_NAME), ConfigurationVariantDetails::mapToOptional);
     return shadeTask;
   }
+
+  /**
+   * Create a copy of the input configuration and returns the copied configuration without the excluded dependencies
+   */
+  private Configuration excludeDependencies(Configuration configuration) {
+    Configuration conf = configuration.copyRecursive();
+    _excludedDependencies.forEach(artifact -> {
+      int idx = artifact.indexOf(':');
+      if (idx == -1) {
+        LOG.info("Will exclude all artifacts having the group: " + artifact + " from the shaded jar");
+        conf.exclude(ImmutableMap.of("group", artifact));
+      } else {
+        LOG.info("Will exclude all artifacts having the group and module: " + artifact + " from the shaded jar");
+        conf.exclude(ImmutableMap.of("group", artifact.substring(0, idx), "module", artifact.substring(idx + 1)));
+      }
+    });
+    return conf;
+  }
+
+  private String shadePrefix(Project project) {
+    return Joiner.on("_").join(
+        sanitize(project.getGroup().toString()),
+        sanitize(project.getName()),
+        sanitize(project.getVersion().toString()));
+  }
+
+  private static String sanitize(String s) {
+    return s.replaceAll("\\-", "_").replaceAll("\\.", "_");
+  }
+
+  private Set<String> classesInConfiguration(FileCollection files) {
+    FileCollection jars = files.getAsFileTree().filter(file -> file.getName().endsWith(".jar"));
+    LOG.info("Will shade the following jars for project.\n{}", Joiner.on("\n").join(new TreeSet<>(jars.getFiles())));
+    return jars.getFiles().stream().flatMap(ShadedJarPackaging::classesInJar).collect(Collectors.toSet());
+  }
+
+  private static Stream<String> classesInJar(File jar) {
+    Set<String> classes = new HashSet<>();
+    try (JarFile jf = new JarFile(jar)) {
+      jf.stream().map(ZipEntry::getName)
+          .filter(name -> name.endsWith(".class"))
+          .map(name -> name.substring(0, name.lastIndexOf('.')))
+          .forEach(classes::add);
+    } catch (IOException e) {
+      throw new RuntimeException("Error reading from Jar file: " + jar, e);
+    }
+    return classes.stream();
+  }
 }

transportable-udfs-plugin/src/main/java/com/linkedin/transport/plugin/packaging/ThinJarPackaging.java

@@ -42,11 +42,10 @@ public List<TaskProvider<? extends Task>> configurePackagingTasks(Project projec
           task.from(platformSourceSet.getOutput());
           task.from(platformSourceSet.getResources());
         });
-
-    String configuration = ShadowBasePlugin.getCONFIGURATION_NAME();
-    project.getArtifacts().add(configuration, thinJarTask);
+
+    project.getArtifacts().add(ShadowBasePlugin.CONFIGURATION_NAME, thinJarTask);
     AdhocComponentWithVariants java = project.getComponents().withType(AdhocComponentWithVariants.class).getByName("java");
-    java.addVariantsFromConfiguration(project.getConfigurations().getByName(configuration), v -> v.mapToOptional());
+    java.addVariantsFromConfiguration(project.getConfigurations().getByName(ShadowBasePlugin.CONFIGURATION_NAME), v -> v.mapToOptional());
 
     return ImmutableList.of(thinJarTask);
   }