LIVY-866: Optimize Yarn GetApplications Query to prevent additional load on Yarn and Livy

server/src/main/scala/org/apache/livy/utils/SparkYarnApp.scala

@@ -16,6 +16,8 @@
  */
 package org.apache.livy.utils
 
+import java.util
+
 import scala.annotation.tailrec
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -25,6 +27,7 @@ import scala.language.postfixOps
 import scala.util.Try
 import scala.util.control.NonFatal
 
+import org.apache.hadoop.yarn.api.protocolrecords.GetApplicationsRequest
 import org.apache.hadoop.yarn.api.records.{ApplicationId, ApplicationReport, FinalApplicationStatus, YarnApplicationState}
 import org.apache.hadoop.yarn.client.api.YarnClient
 import org.apache.hadoop.yarn.conf.YarnConfiguration
@@ -68,27 +71,40 @@ object SparkYarnApp extends Logging {
 
   private var sessionLeakageCheckInterval: Long = _
 
+  /**
+   * Build a GetApplicationsRequest filtered by Spark application type and tags.
+   * Tags are normalized to lowercase to match YARN's tag storage behavior.
+   */
+  private def createGetApplicationsRequest(appTags: util.Set[String]): GetApplicationsRequest = {
+    val normalizedTags = new util.HashSet[String]()
+    appTags.asScala.foreach(tag => normalizedTags.add(tag.toLowerCase))
+    val request = GetApplicationsRequest.newInstance(appType)
+    request.setApplicationTags(normalizedTags)
+    request
+  }
+
   private val leakedAppsGCThread = new Thread() {
     override def run(): Unit = {
-      val client = {
-        mockYarnClient match {
-          case Some(client) => client
-          case None => yarnClient
-        }
+      val client = mockYarnClient match {
+        case Some(client) => client
+        case None => yarnClient
       }
 
       while (true) {
         if (!leakedAppTags.isEmpty) {
           // kill the app if found it and remove it if exceeding a threshold
           val iter = leakedAppTags.entrySet().iterator()
           val now = System.currentTimeMillis()
-          val apps = client.getApplications(appType).asScala
+          val tagSet = new util.HashSet[String](leakedAppTags.keySet())
+          val request = createGetApplicationsRequest(tagSet)
+          val apps = client.getApplications(request).asScala
 
           while(iter.hasNext) {
             var isRemoved = false
             val entry = iter.next()
+            val tagLowerCase = entry.getKey.toLowerCase()
 
-            apps.find(_.getApplicationTags.contains(entry.getKey))
+            apps.find(_.getApplicationTags.contains(tagLowerCase))
               .foreach({ e =>
                 info(s"Kill leaked app ${e.getApplicationId}")
                 client.killApplication(e.getApplicationId)
@@ -196,10 +212,10 @@ class SparkYarnApp private[utils] (
     }
 
     val appTagLowerCase = appTag.toLowerCase()
-
-    // FIXME Should not loop thru all YARN applications but YarnClient doesn't offer an API.
-    // Consider calling rmClient in YarnClient directly.
-    yarnClient.getApplications(appType).asScala.find(_.getApplicationTags.contains(appTagLowerCase))
+    val appTags: util.Set[String] = util.Collections.singleton(appTagLowerCase)
+    val request = createGetApplicationsRequest(appTags)
+    val applicationReports = yarnClient.getApplications(request)
+    applicationReports.asScala.find(_.getApplicationTags.contains(appTagLowerCase))
     match {
       case Some(app) => app.getApplicationId
       case None =>

server/src/test/scala/org/apache/livy/utils/SparkYarnAppSpec.scala

@@ -16,20 +16,22 @@
  */
 package org.apache.livy.utils
 
-import java.util.ArrayList
 import java.util.concurrent.{CountDownLatch, TimeUnit}
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 
 import scala.collection.JavaConverters._
 import scala.concurrent.duration._
 import scala.language.postfixOps
 
+import org.apache.hadoop.yarn.api.protocolrecords.GetApplicationsRequest
 import org.apache.hadoop.yarn.api.records._
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus.UNDEFINED
 import org.apache.hadoop.yarn.api.records.YarnApplicationState._
 import org.apache.hadoop.yarn.client.api.YarnClient
 import org.apache.hadoop.yarn.exceptions.ApplicationAttemptNotFoundException
 import org.apache.hadoop.yarn.util.ConverterUtils
+import org.mockito.ArgumentCaptor
+import org.mockito.Matchers.any
 import org.mockito.Mockito._
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
@@ -49,6 +51,27 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite {
     Thread.`yield`()
   }
 
+  private def mockGetApplicationsByTags(
+      client: YarnClient,
+      reports: List[ApplicationReport]): Unit = {
+    when(client.getApplications(any(classOf[GetApplicationsRequest])))
+      .thenReturn(reports.asJava)
+  }
+
+  private def verifyFilteredGetApplicationsRequest(
+      client: YarnClient,
+      expectedTags: Set[String]): Unit = {
+    val requestCaptor = ArgumentCaptor.forClass(classOf[GetApplicationsRequest])
+    verify(client, atLeastOnce()).getApplications(requestCaptor.capture())
+    val capturedRequest = requestCaptor.getValue
+    expectedTags.foreach { tag =>
+      assert(capturedRequest.getApplicationTags.contains(tag.toLowerCase),
+        s"Request must contain the lowercase tag '$tag'")
+    }
+    assert(capturedRequest.getApplicationTypes.contains("SPARK"),
+      "Request must filter by application type 'SPARK'")
+  }
+
   describe("SparkYarnApp") {
     val TEST_TIMEOUT = 30 seconds
     val appId = ConverterUtils.toApplicationId("application_1467912463905_0021")
@@ -389,8 +412,7 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite {
         when(mockAppReport.getFinalApplicationStatus).thenReturn(FinalApplicationStatus.SUCCEEDED)
         when(mockAppReport.getYarnApplicationState).thenReturn(YarnApplicationState.FINISHED)
         when(mockYarnClient.getApplicationReport(appId)).thenReturn(mockAppReport)
-        when(mockYarnClient.getApplications(Set("SPARK").asJava))
-          .thenReturn(List(mockAppReport).asJava)
+        mockGetApplicationsByTags(mockYarnClient, List(mockAppReport))
 
         val mockListener = mock[SparkAppListener]
         val mockSparkSubmit = mock[LineBufferedProcess]
@@ -404,6 +426,7 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite {
 
           verify(mockYarnClient, atLeast(1)).getApplicationReport(appId)
           verify(mockListener).appIdKnown(appId.toString)
+          verifyFilteredGetApplicationsRequest(mockYarnClient, Set(appTag))
         }
       }
     }
@@ -673,17 +696,18 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite {
         livyConf.set(LivyConf.YARN_APP_LEAKAGE_CHECK_TIMEOUT, "1000ms")
 
         val client = mock[YarnClient]
-        when(client.getApplications(SparkYarnApp.appType)).
-          thenReturn(new ArrayList[ApplicationReport]())
+        mockGetApplicationsByTags(client, List.empty)
 
         SparkYarnApp.init(livyConf, Some(client))
 
         SparkYarnApp.leakedAppTags.clear()
-        SparkYarnApp.leakedAppTags.put("leakApp", System.currentTimeMillis())
+        val leakAppTag = "leakApp"
+        SparkYarnApp.leakedAppTags.put(leakAppTag, System.currentTimeMillis())
 
         Eventually.eventually(Eventually.timeout(TEST_TIMEOUT), Eventually.interval(100 millis)) {
           assert(SparkYarnApp.leakedAppTags.size() == 0)
         }
+        verifyFilteredGetApplicationsRequest(client, Set(leakAppTag))
       }
     }