This is an automated email from the ASF dual-hosted git repository.
abstractdog pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hive.git
The following commit(s) were added to refs/heads/master by this push:
new 3df58f1bef4 HIVE-29780: Fix LLAP plugin rendezvous race under
concurrent DAGs (NullPointerException in LlapTaskCommunicator.setResponse)
(#6659)
3df58f1bef4 is described below
commit 3df58f1bef47b64f24ccf95eb75c67d261438a7f
Author: Bodor Laszlo <[email protected]>
AuthorDate: Wed Aug 26 13:29:34 2026 +0200
HIVE-29780: Fix LLAP plugin rendezvous race under concurrent DAGs
(NullPointerException in LlapTaskCommunicator.setResponse) (#6659)
---
.../hive/llap/tezplugins/LlapPluginBroker.java | 139 +++++++++++
.../hive/llap/tezplugins/LlapTaskCommunicator.java | 28 +--
.../llap/tezplugins/LlapTaskSchedulerService.java | 24 +-
.../hive/llap/tezplugins/TestLlapPluginBroker.java | 262 +++++++++++++++++++++
4 files changed, 421 insertions(+), 32 deletions(-)
diff --git
a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapPluginBroker.java
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapPluginBroker.java
new file mode 100644
index 00000000000..b475d1dfcb2
--- /dev/null
+++
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapPluginBroker.java
@@ -0,0 +1,139 @@
+/*
+ * 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.HashMap;
+import java.util.Map;
+
+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();
+
+ /**
+ * Single source of mutual exclusion for the two maps below: every read,
write, and iteration
+ * happens while this lock is held, so the maps themselves need no internal
synchronization.
+ * Held only for the tiny "look at the peer map, and either pair or park"
window on both sides.
+ */
+ private final Object lock = new Object();
+
+ private final Map<ApplicationAttemptId, LlapTaskCommunicator>
pendingCommunicators =
+ new HashMap<>();
+ private final Map<ApplicationAttemptId, LlapTaskSchedulerService>
pendingSchedulers =
+ new HashMap<>();
+
+ 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();
+ }
+ }
+}
diff --git
a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
index ce5d4dbaccf..c7425df638f 100644
---
a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
+++
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
@@ -84,6 +84,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;
@@ -131,6 +132,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.
@@ -142,13 +144,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);
@@ -163,17 +158,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")
@@ -195,6 +181,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) {
@@ -246,6 +237,7 @@ public void start() {
@Override
public void shutdown() {
super.shutdown();
+ LlapPluginBroker.INSTANCE.unregisterCommunicator(appAttemptId, this);
if (this.communicator != null) {
this.communicator.stop();
}
diff --git
a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskSchedulerService.java
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskSchedulerService.java
index 612e253c5c9..8f4c32a1821 100644
---
a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskSchedulerService.java
+++
b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskSchedulerService.java
@@ -29,6 +29,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;
@@ -203,9 +204,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;
@@ -475,17 +474,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) {
@@ -945,6 +935,7 @@ private void stopTimeoutMonitor() {
@Override
public void shutdown() {
+ LlapPluginBroker.INSTANCE.unregisterScheduler(appAttemptId, this);
writeLock.lock();
try {
if (!this.isStopped.getAndSet(true)) {
@@ -3176,6 +3167,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);
diff --git
a/llap-tez/src/test/org/apache/hadoop/hive/llap/tezplugins/TestLlapPluginBroker.java
b/llap-tez/src/test/org/apache/hadoop/hive/llap/tezplugins/TestLlapPluginBroker.java
new file mode 100644
index 00000000000..19cbd2fa770
--- /dev/null
+++
b/llap-tez/src/test/org/apache/hadoop/hive/llap/tezplugins/TestLlapPluginBroker.java
@@ -0,0 +1,262 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.tez.common.TezUtils;
+import org.apache.tez.dag.api.UserPayload;
+import org.apache.tez.serviceplugins.api.DagInfo;
+import org.apache.tez.serviceplugins.api.TaskCommunicatorContext;
+import org.apache.tez.serviceplugins.api.TaskSchedulerContext;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+/**
+ * Concurrency test for the LLAP plugin broker.
+ *
+ * <p>Fans out N pairs of {@link LlapTaskCommunicator} + {@link
LlapTaskSchedulerService}
+ * constructions in parallel, each pair sharing a unique {@link
ApplicationAttemptId}. All 2N
+ * constructor threads release from a common barrier so their {@link
LlapPluginBroker}
+ * pair-or-park sections interleave as aggressively as the JVM will let them.
After every
+ * construction has completed, the test asserts that every communicator's
paired {@code scheduler}
+ * field refers to the scheduler for the <i>same</i> {@code
ApplicationAttemptId} — i.e. no
+ * cross-DAG mis-pairing.
+ *
+ * <p>An earlier version of both plugins used a single class-static {@code
instance} slot for the
+ * cross-plugin handshake and routinely mis-paired under this workload; the
surfacing symptom in
+ * production was {@code NullPointerException: … scheduler is null} from
+ * {@code LlapTaskCommunicator$3.setResponse}.
+ */
+public class TestLlapPluginBroker {
+
+ private static final int PAIR_COUNT = 32;
+ private static final int ITERATIONS = 25;
+ private static final long AWAIT_TIMEOUT_SECONDS = 60;
+
+ private ExecutorService executor;
+
+ @Before
+ public void setUp() {
+ executor = Executors.newFixedThreadPool(2 * PAIR_COUNT,
+ new
ThreadFactoryBuilder().setDaemon(true).setNameFormat("plugin-broker-%d").build());
+ }
+
+ @After
+ public void tearDown() {
+ executor.shutdownNow();
+ // Reset broker state between iterations so a leaked entry from one test
method doesn't
+ // pollute the next.
+ LlapPluginBroker.INSTANCE.clear();
+ }
+
+ @Test
+ public void testConcurrentPairsCorrectly() throws Exception {
+ for (int iteration = 0; iteration < ITERATIONS; iteration++) {
+ runOneIteration(iteration);
+ }
+ }
+
+ /**
+ * One iteration of the broker stress: build {@link #PAIR_COUNT} plugin
pairs concurrently
+ * and verify that every communicator ended up bound to the scheduler for
its own
+ * {@link ApplicationAttemptId}. Runs assertions inline so a failure in any
iteration fails the
+ * whole test with a specific pair index.
+ */
+ private void runOneIteration(int iteration) throws Exception {
+ List<ApplicationAttemptId> appAttemptIds = new ArrayList<>(PAIR_COUNT);
+ for (int p = 0; p < PAIR_COUNT; p++) {
+ // A distinct ApplicationId per pair — the (iteration, p) tuple keeps
IDs unique across
+ // iterations too, so a stale entry from a previous iteration can't
accidentally match.
+ ApplicationId appId = ApplicationId.newInstance(1000000L + iteration, p
+ 1);
+ appAttemptIds.add(ApplicationAttemptId.newInstance(appId, 1));
+ }
+
+ // Two synchronization primitives:
+ // * a barrier so all 2N threads reach the plugin ctor together,
maximizing the race window,
+ // * a latch so the main thread waits for every ctor to complete.
+ final CyclicBarrier barrier = new CyclicBarrier(2 * PAIR_COUNT);
+ final CountDownLatch done = new CountDownLatch(2 * PAIR_COUNT);
+
+ List<AtomicReference<LlapTaskCommunicator>> comms = new
ArrayList<>(PAIR_COUNT);
+ List<AtomicReference<LlapTaskSchedulerService>> schedulers = new
ArrayList<>(PAIR_COUNT);
+ List<AtomicReference<Throwable>> errors = Collections.synchronizedList(new
ArrayList<>());
+
+ for (int p = 0; p < PAIR_COUNT; p++) {
+ comms.add(new AtomicReference<>());
+ schedulers.add(new AtomicReference<>());
+ }
+
+ for (int p = 0; p < PAIR_COUNT; p++) {
+ final int idx = p;
+ final ApplicationAttemptId appAttemptId = appAttemptIds.get(p);
+
+ executor.submit(() -> {
+ try {
+ barrier.await();
+ comms.get(idx).set(new
LlapTaskCommunicatorForBrokerTest(mockCommContext(appAttemptId)));
+ } catch (Throwable t) {
+ AtomicReference<Throwable> slot = new AtomicReference<>(t);
+ errors.add(slot);
+ } finally {
+ done.countDown();
+ }
+ });
+
+ executor.submit(() -> {
+ try {
+ barrier.await();
+ schedulers.get(idx).set(new
LlapTaskSchedulerServiceForBrokerTest(mockSchedulerContext(appAttemptId)));
+ } catch (Throwable t) {
+ AtomicReference<Throwable> slot = new AtomicReference<>(t);
+ errors.add(slot);
+ } finally {
+ done.countDown();
+ }
+ });
+ }
+
+ assertTrue("all ctors must complete within " + AWAIT_TIMEOUT_SECONDS + "s",
+ done.await(AWAIT_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+
+ if (!errors.isEmpty()) {
+ Throwable first = errors.get(0).get();
+ throw new AssertionError("iteration " + iteration + " had "
+ + errors.size() + " ctor errors; first: " + first, first);
+ }
+
+ // Each pair must be internally consistent: comm.scheduler == matching
scheduler,
+ // scheduler's paired communicator == matching comm.
+ for (int p = 0; p < PAIR_COUNT; p++) {
+ LlapTaskCommunicator c = comms.get(p).get();
+ LlapTaskSchedulerService s = schedulers.get(p).get();
+ assertSame("iteration " + iteration + " pair " + p + ": comm.scheduler
must be its own peer",
+ s, c.getScheduler());
+ assertSame("iteration " + iteration + " pair " + p + ":
scheduler.communicator must be its own peer",
+ c, s.getTaskCommunicator());
+ }
+
+ // After every pair has been brokered, both maps must be empty — pairing
removes the parked
+ // side, and shutdown() reaps orphans. Leftover entries would indicate
either a cross-pair
+ // mis-pairing (which the assertions above already catch) or a leak.
+ assertTrue("pendingCommunicators must drain",
+ LlapPluginBroker.INSTANCE.pendingCommunicatorsView().isEmpty());
+ assertTrue("pendingSchedulers must drain",
+ LlapPluginBroker.INSTANCE.pendingSchedulersView().isEmpty());
+ }
+
+ /**
+ * Verifies the {@code shutdown()} reap: park a plugin without its peer ever
arriving, then
+ * call {@code shutdown()} — the map must be empty afterwards.
+ */
+ @Test
+ public void testShutdownReapsOrphanedEntry() throws Exception {
+ ApplicationId appId = ApplicationId.newInstance(200000L, 1);
+ ApplicationAttemptId appAttemptId =
ApplicationAttemptId.newInstance(appId, 1);
+
+ LlapTaskCommunicator comm =
+ new LlapTaskCommunicatorForBrokerTest(mockCommContext(appAttemptId));
+ assertSame("comm parks itself when peer scheduler is absent", comm,
+
LlapPluginBroker.INSTANCE.pendingCommunicatorsView().get(appAttemptId));
+
+ comm.shutdown();
+ assertFalse("shutdown() must reap our parked entry",
+
LlapPluginBroker.INSTANCE.pendingCommunicatorsView().containsKey(appAttemptId));
+
+ // Symmetric case for the scheduler.
+ LlapTaskSchedulerService scheduler =
+ new
LlapTaskSchedulerServiceForBrokerTest(mockSchedulerContext(appAttemptId));
+ assertSame("scheduler parks itself when peer comm is absent", scheduler,
+ LlapPluginBroker.INSTANCE.pendingSchedulersView().get(appAttemptId));
+
+ scheduler.shutdown();
+ assertFalse("shutdown() must reap our parked entry",
+
LlapPluginBroker.INSTANCE.pendingSchedulersView().containsKey(appAttemptId));
+ }
+
+ private static TaskCommunicatorContext mockCommContext(ApplicationAttemptId
appAttemptId)
+ throws Exception {
+ TaskCommunicatorContext ctx = mock(TaskCommunicatorContext.class);
+ doReturn(appAttemptId).when(ctx).getApplicationAttemptId();
+ doReturn(new Credentials()).when(ctx).getAMCredentials();
+ Configuration conf = new Configuration(false);
+ HiveConf.setVar(conf, ConfVars.LLAP_DAEMON_SERVICE_HOSTS,
"fake-non-zk-cluster");
+
doReturn(TezUtils.createUserPayloadFromConf(conf)).when(ctx).getInitialUserPayload();
+
doReturn(appAttemptId.getApplicationId().toString()).when(ctx).getCurrentAppIdentifier();
+ doReturn(mock(DagInfo.class)).when(ctx).getCurrentDagInfo();
+ doReturn(new
ArrayList<String>()).when(ctx).getInputVertexNames(org.mockito.ArgumentMatchers.any());
+ return ctx;
+ }
+
+ private static TaskSchedulerContext
mockSchedulerContext(ApplicationAttemptId appAttemptId)
+ throws Exception {
+ TaskSchedulerContext ctx = mock(TaskSchedulerContext.class);
+ doReturn(appAttemptId).when(ctx).getApplicationAttemptId();
+ doReturn(11111L).when(ctx).getCustomClusterIdentifier();
+ Configuration conf = new Configuration(false);
+ HiveConf.setVar(conf, ConfVars.LLAP_DAEMON_SERVICE_HOSTS,
"fake-non-zk-cluster");
+ HiveConf.setVar(conf, ConfVars.LLAP_TASK_SCHEDULER_AM_REGISTRY_NAME, "");
+
doReturn(TezUtils.createUserPayloadFromConf(conf)).when(ctx).getInitialUserPayload();
+ return ctx;
+ }
+
+ /**
+ * Minimal LlapTaskCommunicator subclass that skips RPC server startup so we
can spin up many
+ * of them in a single JVM without port collisions.
+ */
+ private static final class LlapTaskCommunicatorForBrokerTest extends
LlapTaskCommunicator {
+ LlapTaskCommunicatorForBrokerTest(TaskCommunicatorContext ctx) {
+ super(ctx);
+ }
+ @Override
+ protected void startRpcServer() {
+ // no-op — we do not exercise task submission in this test
+ }
+ }
+
+ /**
+ * Minimal LlapTaskSchedulerService subclass — the base ctor does the broker
registration,
+ * which is all we need.
+ */
+ private static final class LlapTaskSchedulerServiceForBrokerTest
+ extends LlapTaskSchedulerService {
+ LlapTaskSchedulerServiceForBrokerTest(TaskSchedulerContext ctx) {
+ super(ctx, new org.apache.hadoop.yarn.util.MonotonicClock(), false);
+ }
+ }
+}
