HIVE-29780: Fix LLAP plugin rendezvous race under concurrent DAGs (NullPointerException in LlapTaskCommunicator.setResponse)

llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapPluginBroker.java

@@ -0,0 +1,140 @@
+/*
+ * Licensed 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.hadoop.hive.llap.tezplugins;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+
+import com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Broker that pairs up the two LLAP-side Tez plugins.
+ *
+ * <p>{@link LlapTaskCommunicator} and {@link LlapTaskSchedulerService} are constructed
+ * independently by the Tez app-master and have to find each other so the communicator can call
+ * {@code scheduler.notifyStarted(taskAttemptID)} from the {@code SubmitWork} response callback.
+ * Tez offers no first-class way for the two to meet, so they do it here: whichever side
+ * initializes first for a given DAG parks itself in the appropriate map, and the second side
+ * finds it and pairs the two up.
+ *
+ * <p>The handshake is keyed on {@link ApplicationAttemptId} — both plugin contexts expose it,
+ * and both plugins for the same DAG share it. Keying on the attempt id confines each handshake
+ * to a single DAG so that, in JVMs hosting concurrent DAGs (MiniHS2, MiniLlapCluster, tests
+ * fanning out concurrent inserts), the communicator for DAG-A cannot accidentally pair with the
+ * scheduler for DAG-B.
+ *
+ * <p>A single JVM-wide {@link #INSTANCE} is shared by all plugin constructions. The pair-or-park
+ * decision is atomic under {@link #lock}, so the pairing side always observes a consistent view
+ * of the peer map. {@code shutdown()} on each plugin unconditionally calls the matching
+ * {@code unregister…} method so a partial DAG init (only one of the two plugins was constructed
+ * before something failed) does not leak an entry — entries are keyed on the never-reused
+ * {@link ApplicationAttemptId} and would otherwise sit in the map until the JVM died.
+ *
+ * <p>Production impact of this class is nil. A real Tez {@code DAGAppMaster} runs at most one DAG
+ * over its lifetime, so both maps hold at most one entry, put and removed at plugin construction
+ * / shutdown. Nothing on the task submission or scheduling hot path touches this class.
+ */
+final class LlapPluginBroker {
+
+  static final LlapPluginBroker INSTANCE = new LlapPluginBroker();
+
+  /**
+   * Guards the atomic "look at the peer map, and either pair or park" section on both sides.
+   * Held for the tiny window of the two-branch check; the maps are also {@link ConcurrentMap}s
+   * so {@code unregister…} outside the pairing window is safe.
+   */
+  private final Object lock = new Object();
+
+  private final ConcurrentMap<ApplicationAttemptId, LlapTaskCommunicator> pendingCommunicators =
+      new ConcurrentHashMap<>();
+  private final ConcurrentMap<ApplicationAttemptId, LlapTaskSchedulerService> pendingSchedulers =
+      new ConcurrentHashMap<>();
+
+  private LlapPluginBroker() { }
+
+  /**
+   * Called from the {@link LlapTaskCommunicator} constructor. If the peer scheduler for this DAG
+   * has already parked itself, wire the two together and remove the parked entry. Otherwise park
+   * this communicator so the scheduler picks it up when it arrives.
+   */
+  void registerCommunicator(ApplicationAttemptId appAttemptId, LlapTaskCommunicator communicator) {
+    synchronized (lock) {
+      LlapTaskSchedulerService peer = pendingSchedulers.remove(appAttemptId);
+      if (peer != null) {
+        // We are the last of the pair to initialize for this DAG.
+        peer.setTaskCommunicator(communicator);
+        communicator.setScheduler(peer);
+      } else {
+        pendingCommunicators.put(appAttemptId, communicator);
+      }
+    }
+  }
+
+  /**
+   * Symmetric to {@link #registerCommunicator}: called from the {@link LlapTaskSchedulerService}
+   * constructor.
+   */
+  void registerScheduler(ApplicationAttemptId appAttemptId, LlapTaskSchedulerService scheduler) {
+    synchronized (lock) {
+      LlapTaskCommunicator peer = pendingCommunicators.remove(appAttemptId);
+      if (peer != null) {
+        // We are the last of the pair to initialize for this DAG.
+        scheduler.setTaskCommunicator(peer);
+        peer.setScheduler(scheduler);
+      } else {
+        pendingSchedulers.put(appAttemptId, scheduler);
+      }
+    }
+  }
+
+  /**
+   * Reap a parked communicator entry on shutdown. Safe to call unconditionally: when the plugin
+   * did pair with a peer, the peer's constructor already removed the entry and the two-arg
+   * {@code remove} here is a no-op.
+   */
+  void unregisterCommunicator(ApplicationAttemptId appAttemptId, LlapTaskCommunicator communicator) {
+    synchronized (lock) {
+      pendingCommunicators.remove(appAttemptId, communicator);
+    }
+  }
+
+  /** Symmetric to {@link #unregisterCommunicator}. */
+  void unregisterScheduler(ApplicationAttemptId appAttemptId, LlapTaskSchedulerService scheduler) {
+    synchronized (lock) {
+      pendingSchedulers.remove(appAttemptId, scheduler);
+    }
+  }
+
+  @VisibleForTesting
+  Map<ApplicationAttemptId, LlapTaskCommunicator> pendingCommunicatorsView() {
+    return pendingCommunicators;
+  }
+
+  @VisibleForTesting
+  Map<ApplicationAttemptId, LlapTaskSchedulerService> pendingSchedulersView() {
+    return pendingSchedulers;
+  }
+
+  /** Drop all parked entries. For test cleanup between iterations. */
+  @VisibleForTesting
+  void clear() {
+    synchronized (lock) {
+      pendingCommunicators.clear();
+      pendingSchedulers.clear();
+    }
+  }
+}

llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java

@@ -79,6 +79,7 @@
 import org.apache.hadoop.security.token.Token;
 import org.apache.hadoop.security.token.SecretManager.InvalidToken;
 import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.NodeId;
@@ -126,6 +127,7 @@ public class LlapTaskCommunicator extends TezTaskCommunicatorImpl {
   private final String user;
   private String amHost;
   private String timelineServerUri;
+  private final ApplicationAttemptId appAttemptId;
 
   // These two structures track the list of known nodes, and the list of nodes which are sending in keep-alive heartbeats.
   // Primarily for debugging purposes a.t.m, since there's some unexplained TASK_TIMEOUTS which are currently being observed.
@@ -137,13 +139,6 @@ public class LlapTaskCommunicator extends TezTaskCommunicatorImpl {
   private volatile QueryIdentifierProto currentQueryIdentifierProto;
   private volatile String currentHiveQueryId;
 
-  // TODO: this is an ugly hack because Tez plugin isolation does not make sense for LLAP plugins.
-  //       We are going to register a thread-local here for now, so that the scheduler, initializing
-  //       in the same thread after the communicator, will pick up. Or the other way around.
-  //       This only lives for the duration of the service init.
-  static final Object pluginInitLock = new Object();
-  static LlapTaskCommunicator instance = null;
-
   public LlapTaskCommunicator(
       TaskCommunicatorContext taskCommunicatorContext) {
     super(taskCommunicatorContext);
@@ -158,17 +153,8 @@ public LlapTaskCommunicator(
 
     credentialMap = new ConcurrentHashMap<>();
     sourceStateTracker = new SourceStateTracker(getContext(), this);
-    synchronized (pluginInitLock) {
-      LlapTaskSchedulerService peer = LlapTaskSchedulerService.instance;
-      if (peer != null) {
-        // We are the last to initialize.
-        peer.setTaskCommunicator(this);
-        this.setScheduler(peer);
-        LlapTaskSchedulerService.instance = null;
-      } else {
-        instance = this;
-      }
-    }
+    this.appAttemptId = getContext().getApplicationAttemptId();
+    LlapPluginBroker.INSTANCE.registerCommunicator(appAttemptId, this);
   }
 
   @SuppressWarnings("unchecked")
@@ -190,6 +176,11 @@ void setScheduler(LlapTaskSchedulerService peer) {
     this.scheduler = peer;
   }
 
+  @VisibleForTesting
+  LlapTaskSchedulerService getScheduler() {
+    return scheduler;
+  }
+
   private static final String LLAP_TOKEN_NAME = LlapTokenIdentifier.KIND_NAME.toString();
 
   private void processSendError(Throwable t) {
@@ -241,6 +232,7 @@ public void start() {
   @Override
   public void shutdown() {
     super.shutdown();
+    LlapPluginBroker.INSTANCE.unregisterCommunicator(appAttemptId, this);
     if (this.communicator != null) {
       this.communicator.stop();
     }

llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskSchedulerService.java

@@ -24,6 +24,7 @@
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.metrics2.MetricsSource;
 import org.apache.hadoop.metrics2.MetricsSystem;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 
 import org.apache.hadoop.hive.registry.impl.TezAmRegistryImpl;
@@ -198,9 +199,7 @@ public void setError(Void v, Throwable t) {
                 "Lock metrics for R/W locks LLAP task scheduler", LOCK_METRICS);
   }
 
-  // TODO: this is an ugly hack; see the same in LlapTaskCommunicator for discussion.
-  //       This only lives for the duration of the service init.
-  static LlapTaskSchedulerService instance = null;
+  private final ApplicationAttemptId appAttemptId;
 
   private final Configuration conf;
 
@@ -470,17 +469,8 @@ public LlapTaskSchedulerService(TaskSchedulerContext taskSchedulerContext, Clock
     this.workloadManagementEnabled =
         !StringUtils.isEmpty(conf.get(ConfVars.HIVE_SERVER2_TEZ_INTERACTIVE_QUEUE.varname, "").trim());
 
-    synchronized (LlapTaskCommunicator.pluginInitLock) {
-      LlapTaskCommunicator peer = LlapTaskCommunicator.instance;
-      if (peer != null) {
-        // We are the last to initialize.
-        this.setTaskCommunicator(peer);
-        peer.setScheduler(this);
-        LlapTaskCommunicator.instance = null;
-      } else {
-        instance = this;
-      }
-    }
+    this.appAttemptId = getContext().getApplicationAttemptId();
+    LlapPluginBroker.INSTANCE.registerScheduler(appAttemptId, this);
   }
 
   private Map<Integer, Set<Integer>> getDependencyInfo(TezDAGID depsDagId) {
@@ -940,6 +930,7 @@ private void stopTimeoutMonitor() {
 
   @Override
   public void shutdown() {
+    LlapPluginBroker.INSTANCE.unregisterScheduler(appAttemptId, this);
     writeLock.lock();
     try {
       if (!this.isStopped.getAndSet(true)) {
@@ -3171,6 +3162,11 @@ void setTaskCommunicator(LlapTaskCommunicator communicator) {
     this.communicator = communicator;
   }
 
+  @VisibleForTesting
+  LlapTaskCommunicator getTaskCommunicator() {
+    return communicator;
+  }
+
 
   protected void sendUpdateMessageAsync(TaskInfo ti, boolean newState) {
     WM_LOG.info("Sending message to " + ti.attemptId + ": " + newState);