[twill] 02/02: (TWILL-271) Add custom ClassLoader dependencies

[ch...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

chtyim pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/twill.git

commit 1acf1b4eddc2687ddf416ac26df852f6e65b931d
Author: Terence Yim <te...@google.com>
AuthorDate: Wed Jan 15 11:31:21 2020 -0800

    (TWILL-271) Add custom ClassLoader dependencies
    
    Signed-off-by: Terence Yim <te...@google.com>
---
 .../org/apache/twill/yarn/YarnTwillPreparer.java   | 24 ++++++++++++++--------
 .../apache/twill/yarn/LogLevelChangeTestRun.java   |  2 +-
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/twill-yarn/src/main/java/org/apache/twill/yarn/YarnTwillPreparer.java b/twill-yarn/src/main/java/org/apache/twill/yarn/YarnTwillPreparer.java
index 1ea2ab3..c69b009 100644
--- a/twill-yarn/src/main/java/org/apache/twill/yarn/YarnTwillPreparer.java
+++ b/twill-yarn/src/main/java/org/apache/twill/yarn/YarnTwillPreparer.java
@@ -113,6 +113,7 @@ import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 import java.util.jar.JarEntry;
 import java.util.jar.JarOutputStream;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -121,12 +122,6 @@ import javax.annotation.Nullable;
 final class YarnTwillPreparer implements TwillPreparer {
 
   private static final Logger LOG = LoggerFactory.getLogger(YarnTwillPreparer.class);
-  private static final Function<Class<?>, String> CLASS_TO_NAME = new Function<Class<?>, String>() {
-    @Override
-    public String apply(Class<?> cls) {
-      return cls.getName();
-    }
-  };
 
   private final Configuration config;
   private final TwillSpecification twillSpec;
@@ -575,13 +570,24 @@ final class YarnTwillPreparer implements TwillPreparer {
 
       // Add the TwillRunnableEventHandler class
       if (twillSpec.getEventHandler() != null) {
-        classes.add(getClassLoader().loadClass(twillSpec.getEventHandler().getClassName()));
+        classes.add(classLoader.loadClass(twillSpec.getEventHandler().getClassName()));
+      }
+
+      // Optionally add the custom classloader class
+      if (classLoaderClassName != null) {
+        try {
+          classes.add(classLoader.loadClass(classLoaderClassName));
+        } catch (ClassNotFoundException e) {
+          // Don't throw if the classloader class is not found, as it can be available
+          // in the target cluster with appropriate classpath setting
+          LOG.debug("Cannot load custom classloader class '{}' when preparing for application launch",
+                    classLoaderClassName);
+        }
       }
 
       // The location name is computed from the MD5 of all the classes names
       // The localized name is always APPLICATION_JAR
-      List<String> classList = Lists.newArrayList(Iterables.transform(classes, CLASS_TO_NAME));
-      Collections.sort(classList);
+      List<String> classList = classes.stream().map(Class::getName).sorted().collect(Collectors.toList());
       Hasher hasher = Hashing.md5().newHasher();
       for (String name : classList) {
         hasher.putString(name);
diff --git a/twill-yarn/src/test/java/org/apache/twill/yarn/LogLevelChangeTestRun.java b/twill-yarn/src/test/java/org/apache/twill/yarn/LogLevelChangeTestRun.java
index a1d8ae6..43787f0 100644
--- a/twill-yarn/src/test/java/org/apache/twill/yarn/LogLevelChangeTestRun.java
+++ b/twill-yarn/src/test/java/org/apache/twill/yarn/LogLevelChangeTestRun.java
@@ -277,7 +277,7 @@ public class LogLevelChangeTestRun extends BaseYarnTest {
       if (matchCount == expectedInstances) {
         return;
       }
-      TimeUnit.MILLISECONDS.sleep(100);
+      TimeUnit.SECONDS.sleep(1);
     }
 
     Assert.fail("Timeout waiting for expected log levels");