[KYUUBI #2690] Make ProcessBuilder.commands immutable

### _Why are the changes needed?_
close #2690

When `SparkProcessBuilder` is constructed, the commands are generated according to `conf` and are immutable, so the tags will not take effect if they are set later.

- Make several existing `ProcessBuilder` commands immutable.
- Set tags when building commands.
- Fix Flink tags may not be set to the right value.

### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [x] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #2691 from cxzl25/KYUUBI-2690.

Closes #2690

863e942 [sychen] make ProcBuilder.commands immutable
341796f [sychen] fix UT
9ca7da9 [sychen] use val commands
5ed10d4 [sychen] use TAG_KEY
f4a830f [sychen] add construct for test
a424b3d [sychen] SparkProcessBuilder commands immutable
70d0d7d [sychen] SparkProcessBuilder commands mutable

Authored-by: sychen <sychen@ctrip.com>
Signed-off-by: ulysses-you <ulyssesyou@apache.org>

Author: cxzl25

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala

@@ -167,21 +167,15 @@ private[kyuubi] class EngineRef(
     builder = engineType match {
       case SPARK_SQL =>
         conf.setIfMissing(SparkProcessBuilder.APP_KEY, defaultEngineName)
-        new SparkProcessBuilder(appUser, conf, extraEngineLog)
+        new SparkProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
       case FLINK_SQL =>
         conf.setIfMissing(FlinkProcessBuilder.APP_KEY, defaultEngineName)
-        new FlinkProcessBuilder(appUser, conf, extraEngineLog)
+        new FlinkProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
       case TRINO =>
-        new TrinoProcessBuilder(appUser, conf, extraEngineLog)
+        new TrinoProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
       case HIVE_SQL =>
-        new HiveProcessBuilder(appUser, conf, extraEngineLog)
+        new HiveProcessBuilder(appUser, conf, engineRefId, extraEngineLog)
     }
-    // TODO: Better to do this inside ProcBuilder
-    KyuubiApplicationManager.tagApplication(
-      engineRefId,
-      builder.shortName,
-      builder.clusterManager(),
-      builder.conf)
 
     MetricsSystem.tracing(_.incCount(ENGINE_TOTAL))
     try {

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KyuubiApplicationManager.scala

@@ -23,6 +23,8 @@ import scala.collection.JavaConverters._
 import scala.util.control.NonFatal
 
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.engine.flink.FlinkProcessBuilder
+import org.apache.kyuubi.engine.spark.SparkProcessBuilder
 import org.apache.kyuubi.service.AbstractService
 
 class KyuubiApplicationManager extends AbstractService("KyuubiApplicationManager") {
@@ -88,20 +90,19 @@ class KyuubiApplicationManager extends AbstractService("KyuubiApplicationManager
 
 object KyuubiApplicationManager {
   private def setupSparkYarnTag(tag: String, conf: KyuubiConf): Unit = {
-    val originalTag = conf.getOption("spark.yarn.tags").map(_ + ",").getOrElse("")
-    val newTag = s"${originalTag}KYUUBI,$tag"
-    conf.set("spark.yarn.tags", newTag)
+    val originalTag = conf.getOption(SparkProcessBuilder.TAG_KEY).map(_ + ",").getOrElse("")
+    val newTag = s"${originalTag}KYUUBI" + Some(tag).filterNot(_.isEmpty).map("," + _).getOrElse("")
+    conf.set(SparkProcessBuilder.TAG_KEY, newTag)
   }
 
   private def setupSparkK8sTag(tag: String, conf: KyuubiConf): Unit = {
     conf.set("spark.kubernetes.driver.label.kyuubi_unique_tag", tag)
   }
 
   private def setupFlinkK8sTag(tag: String, conf: KyuubiConf): Unit = {
-    // TODO: yarn.tags or flink.yarn.tags, the mess of flink settings confuses me now.
-    val originalTag = conf.getOption("yarn.tags").map(_ + ",")
-    val newTag = s"${originalTag}KYUUBI,$tag"
-    conf.set("yarn.tags", newTag)
+    val originalTag = conf.getOption(FlinkProcessBuilder.TAG_KEY).map(_ + ",").getOrElse("")
+    val newTag = s"${originalTag}KYUUBI" + Some(tag).filterNot(_.isEmpty).map("," + _).getOrElse("")
+    conf.set(FlinkProcessBuilder.TAG_KEY, newTag)
   }
 
   /**

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/ProcBuilder.scala

@@ -97,7 +97,7 @@ trait ProcBuilder {
 
   protected def proxyUser: String
 
-  protected def commands: Array[String]
+  protected val commands: Array[String]
 
   def conf: KyuubiConf
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/flink/FlinkProcessBuilder.scala

@@ -24,11 +24,13 @@ import java.util.LinkedHashSet
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 
+import com.google.common.annotations.VisibleForTesting
+
 import org.apache.kyuubi._
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
 import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
-import org.apache.kyuubi.engine.ProcBuilder
+import org.apache.kyuubi.engine.{KyuubiApplicationManager, ProcBuilder}
 import org.apache.kyuubi.operation.log.OperationLog
 
 /**
@@ -37,9 +39,15 @@ import org.apache.kyuubi.operation.log.OperationLog
 class FlinkProcessBuilder(
     override val proxyUser: String,
     override val conf: KyuubiConf,
+    val engineRefId: String,
     val extraEngineLog: Option[OperationLog] = None)
   extends ProcBuilder {
 
+  @VisibleForTesting
+  def this(proxyUser: String, conf: KyuubiConf) {
+    this(proxyUser, conf, "")
+  }
+
   private val flinkHome: String = getEngineHome(shortName)
 
   private val FLINK_HADOOP_CLASSPATH: String = "FLINK_HADOOP_CLASSPATH"
@@ -48,7 +56,8 @@ class FlinkProcessBuilder(
 
   override protected def mainClass: String = "org.apache.kyuubi.engine.flink.FlinkSQLEngine"
 
-  override protected def commands: Array[String] = {
+  override protected val commands: Array[String] = {
+    KyuubiApplicationManager.tagApplication(engineRefId, shortName, clusterManager(), conf)
     val buffer = new ArrayBuffer[String]()
     buffer += executable
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/hive/HiveProcessBuilder.scala

@@ -24,19 +24,27 @@ import java.util.LinkedHashSet
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 
+import com.google.common.annotations.VisibleForTesting
+
 import org.apache.kyuubi._
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.{ENGINE_HIVE_EXTRA_CLASSPATH, ENGINE_HIVE_JAVA_OPTIONS, ENGINE_HIVE_MEMORY}
 import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
-import org.apache.kyuubi.engine.ProcBuilder
+import org.apache.kyuubi.engine.{KyuubiApplicationManager, ProcBuilder}
 import org.apache.kyuubi.operation.log.OperationLog
 
 class HiveProcessBuilder(
     override val proxyUser: String,
     override val conf: KyuubiConf,
+    val engineRefId: String,
     val extraEngineLog: Option[OperationLog] = None)
   extends ProcBuilder with Logging {
 
+  @VisibleForTesting
+  def this(proxyUser: String, conf: KyuubiConf) {
+    this(proxyUser, conf, "")
+  }
+
   private val hiveHome: String = getEngineHome(shortName)
 
   private val HIVE_HADOOP_CLASSPATH: String = "HIVE_HADOOP_CLASSPATH"
@@ -45,7 +53,8 @@ class HiveProcessBuilder(
 
   override protected def mainClass: String = "org.apache.kyuubi.engine.hive.HiveSQLEngine"
 
-  override protected def commands: Array[String] = {
+  override protected val commands: Array[String] = {
+    KyuubiApplicationManager.tagApplication(engineRefId, shortName, clusterManager(), conf)
     val buffer = new ArrayBuffer[String]()
     buffer += executable
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkBatchProcessBuilder.scala

@@ -30,7 +30,7 @@ class SparkBatchProcessBuilder(
     batchId: String,
     batchRequest: BatchRequest,
     override val extraEngineLog: Option[OperationLog])
-  extends SparkProcessBuilder(proxyUser, conf, extraEngineLog) {
+  extends SparkProcessBuilder(proxyUser, conf, batchId, extraEngineLog) {
   import SparkProcessBuilder._
 
   override def mainClass: String = batchRequest.getClassName

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala

@@ -22,21 +22,28 @@ import java.nio.file.Paths
 
 import scala.collection.mutable.ArrayBuffer
 
+import com.google.common.annotations.VisibleForTesting
 import org.apache.hadoop.security.UserGroupInformation
 
 import org.apache.kyuubi._
 import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.engine.ProcBuilder
+import org.apache.kyuubi.engine.{KyuubiApplicationManager, ProcBuilder}
 import org.apache.kyuubi.ha.HighAvailabilityConf
 import org.apache.kyuubi.ha.client.AuthTypes
 import org.apache.kyuubi.operation.log.OperationLog
 
 class SparkProcessBuilder(
     override val proxyUser: String,
     override val conf: KyuubiConf,
+    val engineRefId: String,
     val extraEngineLog: Option[OperationLog] = None)
   extends ProcBuilder with Logging {
 
+  @VisibleForTesting
+  def this(proxyUser: String, conf: KyuubiConf) {
+    this(proxyUser, conf, "")
+  }
+
   import SparkProcessBuilder._
 
   private val sparkHome = getEngineHome(shortName)
@@ -48,6 +55,7 @@ class SparkProcessBuilder(
   override def mainClass: String = "org.apache.kyuubi.engine.spark.SparkSQLEngine"
 
   override protected val commands: Array[String] = {
+    KyuubiApplicationManager.tagApplication(engineRefId, shortName, clusterManager(), conf)
     val buffer = new ArrayBuffer[String]()
     buffer += executable
     buffer += CLASS

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/trino/TrinoProcessBuilder.scala

@@ -24,23 +24,33 @@ import java.util.LinkedHashSet
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 
+import com.google.common.annotations.VisibleForTesting
+
 import org.apache.kyuubi.{Logging, SCALA_COMPILE_VERSION, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.{ENGINE_TRINO_CONNECTION_CATALOG, ENGINE_TRINO_CONNECTION_URL, ENGINE_TRINO_EXTRA_CLASSPATH, ENGINE_TRINO_JAVA_OPTIONS, ENGINE_TRINO_MEMORY}
 import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
-import org.apache.kyuubi.engine.ProcBuilder
+import org.apache.kyuubi.engine.{KyuubiApplicationManager, ProcBuilder}
 import org.apache.kyuubi.operation.log.OperationLog
 
 class TrinoProcessBuilder(
     override val proxyUser: String,
     override val conf: KyuubiConf,
-    val extraEngineLog: Option[OperationLog] = None) extends ProcBuilder with Logging {
+    val engineRefId: String,
+    val extraEngineLog: Option[OperationLog] = None)
+  extends ProcBuilder with Logging {
+
+  @VisibleForTesting
+  def this(proxyUser: String, conf: KyuubiConf) {
+    this(proxyUser, conf, "")
+  }
 
   override protected def module: String = "kyuubi-trino-engine"
 
   override protected def mainClass: String = "org.apache.kyuubi.engine.trino.TrinoSqlEngine"
 
-  override protected def commands: Array[String] = {
+  override protected val commands: Array[String] = {
+    KyuubiApplicationManager.tagApplication(engineRefId, shortName, clusterManager(), conf)
     require(
       conf.get(ENGINE_TRINO_CONNECTION_URL).nonEmpty,
       s"Trino server url can not be null! Please set ${ENGINE_TRINO_CONNECTION_URL.key}")

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/flink/FlinkProcessBuilderSuite.scala

@@ -42,7 +42,9 @@ class FlinkProcessBuilderSuite extends KyuubiFunSuite {
   private def envWithAllHadoop: ListMap[String, String] = envWithoutHadoopCLASSPATH +
     ("FLINK_HADOOP_CLASSPATH" -> s"${File.separator}hadoop")
   private def confStr: String = {
-    conf.getAll.map { case (k, v) => s"\\\n\t--conf $k=$v" }.mkString(" ")
+    conf.clone.set("yarn.tags", "KYUUBI").getAll.map {
+      case (k, v) => s"\\\n\t--conf $k=$v"
+    }.mkString(" ")
   }
   private def compareActualAndExpected(builder: FlinkProcessBuilder) = {
     val actualCommands = builder.toString
@@ -53,7 +55,7 @@ class FlinkProcessBuilderSuite extends KyuubiFunSuite {
       s"$confStr"
     info(s"\n\n actualCommands $actualCommands")
     info(s"\n\n expectedCommands $expectedCommands")
-    assert(actualCommands.equals(expectedCommands))
+    assert(actualCommands == expectedCommands)
   }
 
   private def constructClasspathStr(builder: FlinkProcessBuilder) = {
@@ -100,10 +102,9 @@ class FlinkProcessBuilderSuite extends KyuubiFunSuite {
   }
 
   test("all hadoop related environment variables are configured except FLINK_HADOOP_CLASSPATH") {
-    val builder = new FlinkProcessBuilder("vinoyang", conf) {
+    assertThrows[KyuubiException](new FlinkProcessBuilder("vinoyang", conf) {
       override def env: Map[String, String] = envWithoutHadoopCLASSPATH
-    }
-    assertThrows[KyuubiException](builder.toString)
+    })
   }
 
   test("only FLINK_HADOOP_CLASSPATH environment variables are configured") {

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala

@@ -20,6 +20,7 @@ package org.apache.kyuubi.engine.spark
 import java.io.File
 import java.nio.file.{Files, Path, Paths, StandardOpenOption}
 import java.time.Duration
+import java.util.UUID
 import java.util.concurrent.{Executors, TimeUnit}
 
 import org.scalatest.time.SpanSugar._
@@ -260,7 +261,17 @@ class SparkProcessBuilderSuite extends KerberizedTestHelper with MockitoSugar {
 
     val b1 = new SparkProcessBuilder("test", conf)
     assert(b1.toString.contains(s"--conf spark.files=$testKeytab"))
+  }
 
+  test("SparkProcessBuilder commands immutable") {
+    val conf = KyuubiConf(false)
+    val engineRefId = UUID.randomUUID().toString
+    val pb = new SparkProcessBuilder("", conf, engineRefId)
+    assert(pb.toString.contains(engineRefId))
+    val engineRefId2 = UUID.randomUUID().toString
+    conf.set("spark.yarn.tags", engineRefId2)
+    assert(!pb.toString.contains(engineRefId2))
+    assert(pb.toString.contains(engineRefId))
   }
 }