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 000000000000..525ab8109543 --- /dev/null +++ b/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. + * + *

{@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. + * + *

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. + * + *

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. + * + *

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 pendingCommunicators = + new ConcurrentHashMap<>(); + private final ConcurrentMap 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 pendingCommunicatorsView() { + return pendingCommunicators; + } + + @VisibleForTesting + Map 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 0451848f606c..0da1715b4ae6 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 @@ -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(); } 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 c75800c5546a..41aa4a347242 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 @@ -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> 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); 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 000000000000..19cbd2fa770a --- /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. + * + *

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 same {@code ApplicationAttemptId} — i.e. no + * cross-DAG mis-pairing. + * + *

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 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> comms = new ArrayList<>(PAIR_COUNT); + List> schedulers = new ArrayList<>(PAIR_COUNT); + List> 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 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 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()).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); + } + } +}