(TWILL-227) Disabling caching of FileSystem instance when getting del…

…egation token

- Allows getting delegation token for different users without leaking
memory.
  - The FileSystem.get() by default will cache all FileSystem instances
until end of process.

This closes #46 on Github.

Signed-off-by: Terence Yim <chtyim@apache.org>

twill-yarn/src/main/java/org/apache/twill/internal/yarn/YarnUtils.java

@@ -159,19 +159,18 @@ public static List<Token<?>> addDelegationTokens(Configuration config,
       return ImmutableList.of();
     }
 
-    FileSystem fileSystem = getFileSystem(locationFactory, config);
-
-    if (fileSystem == null) {
-      LOG.warn("Unexpected: LocationFactory is not backed by FileContextLocationFactory");
-      return ImmutableList.of();
-    }
+    try (FileSystem fileSystem = getFileSystem(locationFactory)) {
+      if (fileSystem == null) {
+        return ImmutableList.of();
+      }
 
-    String renewer = YarnUtils.getYarnTokenRenewer(config);
+      String renewer = YarnUtils.getYarnTokenRenewer(config);
 
-    Token<?>[] tokens = fileSystem.addDelegationTokens(renewer, credentials);
-    LOG.debug("Added HDFS DelegationTokens: {}", Arrays.toString(tokens));
+      Token<?>[] tokens = fileSystem.addDelegationTokens(renewer, credentials);
+      LOG.debug("Added HDFS DelegationTokens: {}", Arrays.toString(tokens));
 
-    return tokens == null ? ImmutableList.<Token<?>>of() : ImmutableList.copyOf(tokens);
+      return tokens == null ? ImmutableList.<Token<?>>of() : ImmutableList.copyOf(tokens);
+    }
   }
 
   /**
@@ -318,15 +317,22 @@ private static YarnLocalResource setLocalResourceType(YarnLocalResource localRes
    *         {@code null} will be returned if unable to determine the {@link FileSystem}.
    */
   @Nullable
-  private static FileSystem getFileSystem(LocationFactory locationFactory, Configuration config) throws IOException {
+  private static FileSystem getFileSystem(LocationFactory locationFactory) throws IOException {
     if (locationFactory instanceof ForwardingLocationFactory) {
-      return getFileSystem(((ForwardingLocationFactory) locationFactory).getDelegate(), config);
+      return getFileSystem(((ForwardingLocationFactory) locationFactory).getDelegate());
     }
     // Due to HDFS-10296, for encrypted file systems, FileContext does not acquire the KMS delegation token
     // Since we know we are in Yarn, it is safe to get the FileSystem directly, bypassing LocationFactory.
     if (locationFactory instanceof FileContextLocationFactory) {
+      // Disable caching of FileSystem object, as the FileSystem object is only used to get delegation token for the
+      // current user. Caching it may causes leaking of FileSystem object if the method is called with different users.
+      Configuration config = new Configuration(((FileContextLocationFactory) locationFactory).getConfiguration());
+      String scheme = FileSystem.getDefaultUri(config).getScheme();
+      config.set(String.format("fs.%s.impl.disable.cache", scheme), "true");
       return FileSystem.get(config);
     }
+
+    LOG.warn("Unexpected: LocationFactory is not backed by FileContextLocationFactory");
     return null;
   }