[KYUUBI #7644][FLINK] Prevent concurrent access to bootstrap job IDs

externals/kyuubi-flink-sql-engine/src/main/java/org/apache/flink/client/deployment/application/executors/EmbeddedExecutorFactory.java

@@ -25,6 +25,7 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.client.cli.ClientOptions;
 import org.apache.flink.client.deployment.application.EmbeddedJobClient;
@@ -114,6 +115,8 @@ public class EmbeddedExecutorFactory implements PipelineExecutorFactory {
 
   private static Collection<JobID> bootstrapJobIds;
 
+  private static volatile boolean bootstrapJobIdsClaimed;
+
   private static Collection<JobID> submittedJobIds;
 
   /**
@@ -198,12 +201,13 @@ public EmbeddedExecutorFactory(
     checkState(EmbeddedExecutorFactory.dispatcherGateway == null);
     checkState(EmbeddedExecutorFactory.retryExecutor == null);
     synchronized (bootstrapLock) {
-      // submittedJobIds would be always 1, because we create a new list to avoid concurrent access
-      // issues
+      // Keep Flink's collection for the application bootstrap job. Later Kyuubi jobs use the
+      // thread-safe copy to avoid concurrent access to Flink's ArrayList.
       LOGGER.debug("Bootstrapping EmbeddedExecutorFactory.");
       EmbeddedExecutorFactory.submittedJobIds =
           new ConcurrentLinkedQueue<>(checkNotNull(applicationJobIds));
       EmbeddedExecutorFactory.bootstrapJobIds = applicationJobIds;
+      EmbeddedExecutorFactory.bootstrapJobIdsClaimed = !applicationJobIds.isEmpty();
       EmbeddedExecutorFactory.suspendedJobIds = suspendedJobIds;
       EmbeddedExecutorFactory.terminalJobIds = terminalJobIds;
       EmbeddedExecutorFactory.dispatcherGateway = checkNotNull(dispatcherGateway);
@@ -228,7 +232,20 @@ public boolean isCompatibleWith(final Configuration configuration) {
   @Override
   public PipelineExecutor getExecutor(final Configuration configuration) {
     checkNotNull(configuration);
-    Collection<JobID> executorJobIDs;
+    final Collection<JobID> executorJobIDs = claimJobIdsForExecutor();
+    final EmbeddedJobClientCreator jobClientCreator =
+        (jobId, userCodeClassloader) ->
+            newEmbeddedJobClient(
+                jobId, configuration.get(ClientOptions.CLIENT_TIMEOUT), userCodeClassloader);
+    return stampApplicationId(newEmbeddedExecutor(executorJobIDs, configuration, jobClientCreator));
+  }
+
+  @VisibleForTesting
+  static Collection<JobID> claimJobIdsForExecutor() {
+    if (bootstrapJobIdsClaimed) {
+      LOGGER.info("Submitting new Kyuubi job. Job submitted: {}.", submittedJobIds.size());
+      return submittedJobIds;
+    }
     synchronized (bootstrapLock) {
       // wait in a loop to avoid spurious wakeups
       int retry = 0;
@@ -247,19 +264,17 @@ public PipelineExecutor getExecutor(final Configuration configuration) {
                 + BOOTSTRAP_WAIT_INTERVAL * BOOTSTRAP_WAIT_RETRIES
                 + " ms. Please check the engine log for more details.");
       }
-    }
-    if (bootstrapJobIds.size() > 0) {
+      if (!bootstrapJobIdsClaimed) {
+        // Flink owns this collection and expects the application bootstrap job in it. Claim it
+        // before returning the executor so another submission cannot observe the list as empty and
+        // concurrently add to Flink's non-thread-safe ArrayList.
+        bootstrapJobIdsClaimed = true;
+        LOGGER.info("Bootstrapping Flink SQL engine with the initial SQL.");
+        return bootstrapJobIds;
+      }
       LOGGER.info("Submitting new Kyuubi job. Job submitted: {}.", submittedJobIds.size());
-      executorJobIDs = submittedJobIds;
-    } else {
-      LOGGER.info("Bootstrapping Flink SQL engine with the initial SQL.");
-      executorJobIDs = bootstrapJobIds;
+      return submittedJobIds;
     }
-    final EmbeddedJobClientCreator jobClientCreator =
-        (jobId, userCodeClassloader) ->
-            newEmbeddedJobClient(
-                jobId, configuration.get(ClientOptions.CLIENT_TIMEOUT), userCodeClassloader);
-    return stampApplicationId(newEmbeddedExecutor(executorJobIDs, configuration, jobClientCreator));
   }
 
   private static PipelineExecutor newEmbeddedExecutor(

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/FlinkSQLEngine.scala

@@ -106,11 +106,13 @@ object FlinkSQLEngine extends Logging {
       }
 
       val engineContext = FlinkEngineUtils.getDefaultContext(args, flinkConf, flinkConfDir)
+      // Finish engine-level initialization before exposing the frontend so client jobs cannot race
+      // with the application bootstrap job.
+      bootstrap(executionTarget)
+
       startEngine(engineContext)
       info("Flink engine started")
 
-      bootstrap(executionTarget)
-
       // blocking main thread
       countDownLatch.await()
     } catch {

externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/flink/client/deployment/application/executors/EmbeddedExecutorFactorySuite.scala

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.client.deployment.application.executors
+
+import java.lang.reflect.{InvocationHandler, Method, Proxy}
+import java.util.{ArrayList, Collection}
+
+import org.apache.flink.api.common.JobID
+import org.apache.flink.runtime.dispatcher.DispatcherGateway
+import org.apache.flink.util.concurrent.ScheduledExecutor
+
+import org.apache.kyuubi.KyuubiFunSuite
+
+class EmbeddedExecutorFactorySuite extends KyuubiFunSuite {
+
+  test("reserve Flink application job ids for only one executor") {
+    val applicationJobIds = new ArrayList[JobID]()
+    new EmbeddedExecutorFactory(
+      applicationJobIds,
+      proxy(classOf[DispatcherGateway]),
+      proxy(classOf[ScheduledExecutor]))
+
+    val bootstrapExecutorJobIds = claimJobIdsForExecutor()
+    val statementExecutorJobIds = claimJobIdsForExecutor()
+
+    assert(bootstrapExecutorJobIds eq applicationJobIds)
+    assert(statementExecutorJobIds ne applicationJobIds)
+
+    statementExecutorJobIds.add(new JobID())
+    assert(applicationJobIds.isEmpty)
+  }
+
+  private def claimJobIdsForExecutor(): Collection[JobID] = {
+    val claimMethod = classOf[EmbeddedExecutorFactory]
+      .getDeclaredMethod("claimJobIdsForExecutor")
+    claimMethod.setAccessible(true)
+    claimMethod.invoke(null).asInstanceOf[Collection[JobID]]
+  }
+
+  private def proxy[T](interfaceClass: Class[T]): T = {
+    val invocationHandler = new InvocationHandler {
+      override def invoke(proxy: Object, method: Method, args: Array[Object]): Object = null
+    }
+    Proxy.newProxyInstance(
+      interfaceClass.getClassLoader,
+      Array(interfaceClass),
+      invocationHandler).asInstanceOf[T]
+  }
+}