Write tokens to HDFS and read them back when required, rather than sending them over the wire.

Author: harishreedharan

bin/utils.sh

@@ -36,7 +36,7 @@ function gatherSparkSubmitOpts() {
       --conf | --repositories | --properties-file | --driver-memory | --driver-java-options | \
       --driver-library-path | --driver-class-path | --executor-memory | --driver-cores | \
       --total-executor-cores | --executor-cores | --queue | --num-executors | --archives | \
-      --proxy-user)
+      --proxy-user | --principal | --keytab)
         if [[ $# -lt 2 ]]; then
           "$SUBMIT_USAGE_FUNCTION"
           exit 1;

core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala

@@ -122,9 +122,9 @@ class SparkHadoopUtil extends Logging {
     UserGroupInformation.loginUserFromKeytab(principalName, keytabFilename)
   }
 
-  def setPrincipalAndKeytabForLogin(principal: String, keytab: String): Unit = {}
+  def setPrincipalAndKeytabForLogin(principal: String, keytab: String): Unit = ???
 
-  private[spark] def scheduleLoginFromKeytab(callback: (SerializableBuffer) => Unit): Unit = {}
+  private[spark] def scheduleLoginFromKeytab(callback: (String) => Unit): Unit = {}
 
   /**
    * Returns a function that can be called to find Hadoop FileSystem bytes read. If

core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala

@@ -372,6 +372,8 @@ object SparkSubmit {
       OptionAssigner(args.files, YARN, CLUSTER, clOption = "--files"),
       OptionAssigner(args.archives, YARN, CLUSTER, clOption = "--archives"),
       OptionAssigner(args.jars, YARN, CLUSTER, clOption = "--addJars"),
+      OptionAssigner(args.principal, YARN, CLUSTER, clOption = "--principal"),
+      OptionAssigner(args.keytab, YARN, CLUSTER, clOption = "--keytab"),
 
       // Other options
       OptionAssigner(args.executorMemory, STANDALONE | MESOS | YARN, ALL_DEPLOY_MODES,

core/src/main/scala/org/apache/spark/deploy/SparkSubmitArguments.scala

@@ -58,6 +58,8 @@ private[spark] class SparkSubmitArguments(args: Seq[String], env: Map[String, St
   var action: SparkSubmitAction = null
   val sparkProperties: HashMap[String, String] = new HashMap[String, String]()
   var proxyUser: String = null
+  var principal: String = null
+  var keytab: String = null
 
   // Standalone cluster mode only
   var supervise: Boolean = false
@@ -410,6 +412,14 @@ private[spark] class SparkSubmitArguments(args: Seq[String], env: Map[String, St
         proxyUser = value
         parse(tail)
 
+      case ("--principal") :: value :: tail =>
+        principal = value
+        parse(tail)
+
+      case ("--keytab") :: value :: tail =>
+        keytab = value
+        parse(tail)
+
       case ("--help" | "-h") :: tail =>
         printUsageAndExit(0)
 

core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala

@@ -17,8 +17,8 @@
 
 package org.apache.spark.executor
 
-import java.net.URL
 import java.io.{ByteArrayInputStream, DataInputStream}
+import java.net.URL
 import java.nio.ByteBuffer
 
 import scala.collection.mutable
@@ -27,6 +27,7 @@ import scala.concurrent.Await
 import akka.actor.{Actor, ActorSelection, Props}
 import akka.pattern.Patterns
 import akka.remote.{RemotingLifecycleEvent, DisassociatedEvent}
+import org.apache.hadoop.fs.{Path, FileSystem}
 import org.apache.hadoop.security.Credentials
 
 import org.apache.spark.{Logging, SecurityManager, SparkConf, SparkEnv}
@@ -109,12 +110,14 @@ private[spark] class CoarseGrainedExecutorBackend(
       context.system.shutdown()
 
     // Add new credentials received from the driver to the current user.
-    case UpdateCredentials(newCredentials) =>
+    case UpdateCredentials(newCredentialsPath) =>
       logInfo("New credentials received from driver, adding the credentials to the current user")
       val credentials = new Credentials()
-      credentials.readTokenStorageStream(
-        new DataInputStream(new ByteArrayInputStream(newCredentials.value.array())))
+      val remoteFs = FileSystem.get(SparkHadoopUtil.get.conf)
+      val inStream = remoteFs.open(new Path(newCredentialsPath))
+      credentials.readTokenStorageStream(inStream)
       SparkHadoopUtil.get.addCurrentUserCredentials(credentials)
+      inStream.close()
   }
 
   override def statusUpdate(taskId: Long, state: TaskState, data: ByteBuffer) {

core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala

@@ -53,7 +53,7 @@ private[spark] object CoarseGrainedClusterMessages {
 
   // When the delegation tokens are about expire, the driver creates new tokens and sends them to
   // the executors via this message.
-  case class UpdateCredentials(newCredentials: SerializableBuffer)
+  case class UpdateCredentials(newCredentialsLocation: String)
     extends CoarseGrainedClusterMessage
 
   object StatusUpdate {

core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala

@@ -27,8 +27,9 @@ import akka.actor._
 import akka.pattern.ask
 import akka.remote.{DisassociatedEvent, RemotingLifecycleEvent}
 
-import org.apache.spark.deploy.SparkHadoopUtil
+
 import org.apache.spark.{ExecutorAllocationClient, Logging, SparkEnv, SparkException, TaskState}
+import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.scheduler._
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
 import org.apache.spark.util.{ActorLogReceive, SerializableBuffer, AkkaUtils, Utils}
@@ -75,11 +76,11 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val actorSyste
   /**
    * Send new credentials to executors. This is the method that is called when the scheduled
    * login completes, so the new credentials can be sent to the executors.
-   * @param credentials
+   * @param credentialsPath
    */
-  def sendNewCredentialsToExecutors(credentials: SerializableBuffer): Unit = {
+  def sendNewCredentialsToExecutors(credentialsPath: String): Unit = {
     // We don't care about the reply, so going to deadLetters is fine.
-    executorDataMap.values.foreach(_.executorActor ! UpdateCredentials(credentials))
+    executorDataMap.values.foreach(_.executorActor ! UpdateCredentials(credentialsPath))
   }
 
   class DriverActor(sparkProperties: Seq[(String, String)]) extends Actor with ActorLogReceive {

yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala

@@ -562,25 +562,22 @@ private[spark] class Client(
   }
 
   def setupCredentials(): Unit = {
-    Option(args.principal) match {
-      case Some(principal) =>
-        Option(args.keytab) match {
-          case Some(keytabPath) =>
-            // Generate a file name that can be used for the keytab file, that does not conflict
-            // with any user file.
-            logInfo("Attempting to login to the Kerberos" +
-              s" using principal: $principal and keytab: $keytabPath")
-            val f = new File(keytabPath)
-            keytabFileName = f.getName + "-" + System.currentTimeMillis()
-            val ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(principal, keytabPath)
-            credentials = ugi.getCredentials
-            loginFromKeytab = true
-            logInfo("Successfully logged into Kerberos.")
-          case None =>
-            throw new SparkException("Keytab must be specified when principal is specified.")
-        }
-      case None =>
-        credentials = UserGroupInformation.getCurrentUser.getCredentials
+    if (args.principal != null) {
+      if (args.keytab == null) {
+        throw new SparkException("Keytab must be specified when principal is specified.")
+      }
+      logInfo("Attempting to login to the Kerberos" +
+        s" using principal: ${args.principal} and keytab: ${args.keytab}")
+      val f = new File(args.keytab)
+      // Generate a file name that can be used for the keytab file, that does not conflict
+      // with any user file.
+      keytabFileName = f.getName + "-" + System.currentTimeMillis()
+      val ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(args.principal, args.keytab)
+      credentials = ugi.getCredentials
+      loginFromKeytab = true
+      logInfo("Successfully logged into Kerberos.")
+    } else {
+      credentials = UserGroupInformation.getCurrentUser.getCredentials
     }
   }
 

yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala

@@ -20,8 +20,7 @@ package org.apache.spark.deploy.yarn
 import java.io._
 import java.net.URI
 import java.nio.ByteBuffer
-import java.util.concurrent.atomic.{AtomicReference, AtomicBoolean}
-import java.util.concurrent.{TimeUnit, ThreadFactory, Executors}
+import java.util.concurrent.{TimeUnit, Executors}
 import java.util.regex.Matcher
 import java.util.regex.Pattern
 
@@ -49,10 +48,10 @@ import org.apache.spark.util.{SerializableBuffer, Utils}
  */
 class YarnSparkHadoopUtil extends SparkHadoopUtil {
 
-  private var keytabFile: Option[String] = None
-  private var loginPrincipal: Option[String] = None
-  private val loggedInViaKeytab = new AtomicBoolean(false)
-  private val loggedInUGI = new AtomicReference[UserGroupInformation](null)
+  private var keytab: String = null
+  private var principal: String = null
+  @volatile private var loggedInViaKeytab = false
+  @volatile private var loggedInUGI: UserGroupInformation = null
 
   override def transferCredentials(source: UserGroupInformation, dest: UserGroupInformation) {
     dest.addCredentials(source.getCredentials())
@@ -94,58 +93,61 @@ class YarnSparkHadoopUtil extends SparkHadoopUtil {
   }
 
   override def setPrincipalAndKeytabForLogin(principal: String, keytab: String): Unit = {
-    loginPrincipal = Option(principal)
-    keytabFile = Option(keytab)
+    this.principal = principal
+    this.keytab = keytab
   }
 
   private[spark] override def scheduleLoginFromKeytab(
-    callback: (SerializableBuffer)  => Unit): Unit = {
-
-    loginPrincipal match {
-      case Some(principal) =>
-        val keytab = keytabFile.get
-        val remoteFs = FileSystem.get(conf)
-        val remoteKeytabPath = new Path(
-          remoteFs.getHomeDirectory, System.getenv("SPARK_STAGING_DIR") + Path.SEPARATOR + keytab)
-        val localFS = FileSystem.getLocal(conf)
-        // At this point, SparkEnv is likely no initialized, so create a dir, put the keytab there.
-        val tempDir = Utils.createTempDir()
-        val localURI = new URI(tempDir.getAbsolutePath + Path.SEPARATOR + keytab)
-        val qualifiedURI = new  URI(localFS.makeQualified(new Path(localURI)).toString)
-        FileUtil.copy(
-          remoteFs, remoteKeytabPath, localFS, new Path(qualifiedURI), false, false, conf)
-        // Get the current credentials, find out when they expire.
-        val creds = UserGroupInformation.getCurrentUser.getCredentials
-        val credStream = new ByteArrayOutputStream()
-        creds.writeTokenStorageToStream(new DataOutputStream(credStream))
-        val in = new DataInputStream(new ByteArrayInputStream(credStream.toByteArray))
-        val tokenIdentifier = new DelegationTokenIdentifier()
-        tokenIdentifier.readFields(in)
-        val timeToRenewal = (0.6 * (tokenIdentifier.getMaxDate - System.currentTimeMillis())).toLong
-        Executors.newSingleThreadScheduledExecutor(new ThreadFactory {
-          override def newThread(r: Runnable): Thread = {
-            val t = new Thread(r)
-            t.setName("Delegation Token Refresh Thread")
-            t.setDaemon(true)
-            t
-          }
-        }).scheduleWithFixedDelay(new Runnable {
-          override def run(): Unit = {
-            if (!loggedInViaKeytab.get()) {
-              loggedInUGI.set(UserGroupInformation.loginUserFromKeytabAndReturnUGI(
-                principal, tempDir.getAbsolutePath + Path.SEPARATOR + keytab))
-              loggedInViaKeytab.set(true)
+    callback: (String) => Unit): Unit = {
+    if (principal != null) {
+      val stagingDir = System.getenv("SPARK_YARN_STAGING_DIR")
+      val remoteFs = FileSystem.get(conf)
+      val remoteKeytabPath = new Path(
+        remoteFs.getHomeDirectory, stagingDir + Path.SEPARATOR + keytab)
+      val localFS = FileSystem.getLocal(conf)
+      // At this point, SparkEnv is likely no initialized, so create a dir, put the keytab there.
+      val tempDir = Utils.createTempDir()
+      Utils.chmod700(tempDir)
+      val localURI = new URI(tempDir.getAbsolutePath + Path.SEPARATOR + keytab)
+      val qualifiedURI = new URI(localFS.makeQualified(new Path(localURI)).toString)
+      FileUtil.copy(
+        remoteFs, remoteKeytabPath, localFS, new Path(qualifiedURI), false, false, conf)
+      // Get the current credentials, find out when they expire.
+      val creds = {
+        if (loggedInUGI == null) {
+          UserGroupInformation.getCurrentUser.getCredentials
+        } else {
+          loggedInUGI.getCredentials
+        }
+      }
+      val credStream = new ByteArrayOutputStream()
+      creds.writeTokenStorageToStream(new DataOutputStream(credStream))
+      val in = new DataInputStream(new ByteArrayInputStream(credStream.toByteArray))
+      val tokenIdentifier = new DelegationTokenIdentifier()
+      tokenIdentifier.readFields(in)
+      val timeToRenewal = (0.6 * (tokenIdentifier.getMaxDate - System.currentTimeMillis())).toLong
+      Executors.newSingleThreadScheduledExecutor(
+        Utils.namedThreadFactory("Delegation Token Refresh Thread")).scheduleWithFixedDelay(
+          new Runnable {
+            override def run(): Unit = {
+              if (!loggedInViaKeytab) {
+                loggedInUGI = UserGroupInformation.loginUserFromKeytabAndReturnUGI(
+                  principal, tempDir.getAbsolutePath + Path.SEPARATOR + keytab)
+                loggedInViaKeytab = true
+              }
+              val nns = getNameNodesToAccess(sparkConf) + remoteKeytabPath
+              val newCredentials = loggedInUGI.getCredentials
+              obtainTokensForNamenodes(nns, conf, newCredentials)
+              val tokenPath = new Path(remoteFs.getHomeDirectory, stagingDir + Path.SEPARATOR +
+                "credentials - " + System.currentTimeMillis())
+              val stream = remoteFs.create(tokenPath, true)
+              // Now write this out via Akka to executors.
+              newCredentials.writeTokenStorageToStream(stream)
+              stream.hflush()
+              stream.close()
+              callback(tokenPath.toString)
             }
-            val nns = getNameNodesToAccess(sparkConf) + remoteKeytabPath
-            val newCredentials = loggedInUGI.get().getCredentials
-            obtainTokensForNamenodes(nns, conf, newCredentials)
-            // Now write this out via Akka to executors.
-            val outputStream = new ByteArrayOutputStream()
-            newCredentials.writeTokenStorageToStream(new DataOutputStream(outputStream))
-            callback(new SerializableBuffer(ByteBuffer.wrap(outputStream.toByteArray)))
-          }
-        }, timeToRenewal, timeToRenewal, TimeUnit.MILLISECONDS)
-      case None =>
+          }, timeToRenewal, timeToRenewal, TimeUnit.MILLISECONDS)
     }
   }