abstractdog commented on code in PR #6501: URL: https://github.com/apache/hive/pull/6501#discussion_r3757053580
########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/YarnQueueMetricsCollector.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Objects; + +/** + * Collects YARN queue resource metrics using a shared cache to reduce ResourceManager load. + * Coordinates with other collectors via QueueMetricsCache to prevent duplicate RM calls. + * + * Executor pool management (sizing, lifecycle, JMX) is fully delegated to + * {@link QueueMetricsRefreshPool}. This class focuses solely on per-query/per-queue + * metrics logic: session registration, refresh scheduling, and cache coordination. + */ +public class YarnQueueMetricsCollector implements QueueMetricsCollector { + private static final Logger LOG = LoggerFactory.getLogger(YarnQueueMetricsCollector.class); + + private final YarnClient yarnClient; + private final String queueName; + private final long refreshIntervalMs; + private final String dagName; + + /** + * Creates a collector for the given queue and query. Non-blocking: metrics are + * fetched asynchronously; the first progress update may show no metrics, subsequent + * updates will once the first fetch completes (within one refreshIntervalMs). + * + * @param yarnClient Live YarnClient from the Tez session + * @param queueName YARN queue this query runs on + * @param refreshIntervalMs How often to poll YARN RM (ms) + * @param dagName Hive query identifier used for logging + * @param hiveConf Unused (kept for API compatibility) + */ + public YarnQueueMetricsCollector(YarnClient yarnClient, String queueName, long refreshIntervalMs, String dagName, + HiveConf hiveConf) { Review Comment: I just realized hiveConf is not use, we can remove it ########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/YarnQueueMetricsCollector.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Objects; + +/** + * Collects YARN queue resource metrics using a shared cache to reduce ResourceManager load. + * Coordinates with other collectors via QueueMetricsCache to prevent duplicate RM calls. + * + * Executor pool management (sizing, lifecycle, JMX) is fully delegated to + * {@link QueueMetricsRefreshPool}. This class focuses solely on per-query/per-queue + * metrics logic: session registration, refresh scheduling, and cache coordination. + */ +public class YarnQueueMetricsCollector implements QueueMetricsCollector { + private static final Logger LOG = LoggerFactory.getLogger(YarnQueueMetricsCollector.class); + + private final YarnClient yarnClient; + private final String queueName; + private final long refreshIntervalMs; + private final String dagName; + + /** + * Creates a collector for the given queue and query. Non-blocking: metrics are + * fetched asynchronously; the first progress update may show no metrics, subsequent + * updates will once the first fetch completes (within one refreshIntervalMs). + * + * @param yarnClient Live YarnClient from the Tez session + * @param queueName YARN queue this query runs on + * @param refreshIntervalMs How often to poll YARN RM (ms) + * @param dagName Hive query identifier used for logging + * @param hiveConf Unused (kept for API compatibility) Review Comment: this is a new feature, no API compat is needed ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/TestYarnQueueMetricsCollector.java: ########## @@ -0,0 +1,587 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.api.records.QueueStatistics; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mockingDetails; +import static org.mockito.Mockito.when; + +/** + * Test cases for YarnQueueMetricsCollector. + */ +public class TestYarnQueueMetricsCollector { + + @Mock + private YarnClient mockYarnClient; + + @Mock + private QueueInfo mockQueueInfo; + + @Mock + private QueueStatistics mockQueueStats; + + private AutoCloseable closeable; + private HiveConf testConf; + + private static final long WAIT_TIMEOUT_MS = 5000; + + @Before + public void setUp() { + closeable = MockitoAnnotations.openMocks(this); + testConf = new HiveConf(); + // Reset the pool manager singleton and cache so each test starts with a clean state. + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + @After + public void tearDown() throws Exception { + if (closeable != null) { + closeable.close(); + } + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + /** + * Helper to create a collector in tests using a default HiveConf (min pool sizes). + */ + private YarnQueueMetricsCollector newCollector(YarnClient yarnClient, String queueName, + long refreshIntervalMs, String queryId) { + return new YarnQueueMetricsCollector(yarnClient, queueName, refreshIntervalMs, queryId, testConf); + } + + /** + * Waits for a snapshot to be available (non-null). + */ + private QueueMetricsSnapshot waitForSnapshot( + YarnQueueMetricsCollector collector, long timeoutMs) { + long startTime = System.currentTimeMillis(); + QueueMetricsSnapshot snapshot; + while ((snapshot = collector.getLatestSnapshot()) == null) { + if (System.currentTimeMillis() - startTime > timeoutMs) { + fail("Snapshot not available after " + timeoutMs + "ms"); + } + Thread.onSpinWait(); // Hint to JVM that this is a spin-wait loop Review Comment: okay, makes sense in general, I would not be in favor of introducing new libraries, but `Awaitility` looks good, and it's already present in hive, reusing it makes sense to me, I would consider moving that dependency to the root pom.xml's dependencyManagement and reuse: https://github.com/apache/hive/blob/1d0bf6981788fe2e7de44aaf44cce84c2f7a964e/iceberg/pom.xml#L208-L212 using Awaitility could totally avoid re-defining functions like `waitForInvocationCount` this could also prevent new sonarqube warnings on busy-wait loops ########## common/src/java/org/apache/hadoop/hive/conf/HiveConf.java: ########## @@ -3940,6 +4006,15 @@ public static enum ConfVars { HIVE_SERVER2_TEZ_QUEUE_ACCESS_CHECK("hive.server2.tez.queue.access.check", false, "Whether to check user access to explicitly specified YARN queues. " + "yarn.resourcemanager.webapp.address must be configured to use this."), + HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL("hive.tez.queue.metrics.refresh.interval", "0s", + new TimeValidator(TimeUnit.SECONDS), + "Interval for refreshing YARN queue resource metrics during Tez query execution. " + + "When set to a positive value (e.g. 10s), displays real-time memory, vCore, capacity " + + "and application metrics for the YARN queue being used. " + + "Set to 0 or negative to disable. Minimum effective value is 1 second."), + HIVE_SERVER2_TEZ_QUEUE_METRICS_REFRESH_THREADS("hive.server2.tez.queue.metrics.refresh.threads", 4, + "Number of threads in the scheduled thread pool for refreshing YARN queue metrics. " + + "This pool is used by HiveServer2 to periodically collect queue resource information from YARN RM. "), Review Comment: I meant rather K8s environments, not different execution engines like spark, mr please remove the "isTez" related parts of HIVE_SERVER2_TEZ_QUEUE_METRICS_REFRESH_THREADS ########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/QueueMetricsState.java: ########## @@ -0,0 +1,287 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.OptionalLong; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Holds all runtime state for YARN queue metrics collection on a single queue. + * One instance exists per active queue name in the JVM, stored in {@link QueueMetricsCache}. + * <p> + * Owns all per-queue logic: session interval registration, refresh task scheduling, + * thundering herd prevention, and circuit breaker. All fields are private — callers + * interact only through methods. + * <p> + * Ownership model: + * <ul> + * <li>{@code intervalCounts}, {@code minRefreshIntervalMs}, {@code activeSessionCount} + * — owned by {@link #registerInterval}/{@link #deregisterInterval} via lock-free atomics</li> + * <li>{@code refreshTask}, {@code taskCurrentRefreshIntervalMs} + * — owned by {@link #ensureTaskScheduled} under {@code synchronized(this)}</li> + * <li>{@code snapshot}, {@code lastWriteTime} — written by the refresh thread, read by + * TezProgressMonitor; {@code volatile} for visibility without synchronization</li> + * </ul> + */ +public class QueueMetricsState { + private static final Logger LOG = LoggerFactory.getLogger(QueueMetricsState.class); + + private static final int MAX_CONSECUTIVE_FAILURES = 5; + private static final int CIRCUIT_BREAKER_PROBE_INTERVAL = 10; + + // Metrics data (written by refresh thread, read by TezProgressMonitor) + private final AtomicReference<QueueMetricsSnapshot> snapshot; + private volatile long lastWriteTime; + + // Session interval tracking (lock-free atomics) + private final AtomicLong minRefreshIntervalMs; + private final ConcurrentHashMap<Long, AtomicInteger> intervalCounts = new ConcurrentHashMap<>(); + private final AtomicInteger activeSessionCount = new AtomicInteger(0); + + // Refresh task (owned by ensureTaskScheduled under synchronized(this)) + private final AtomicReference<ScheduledFuture<?>> refreshTask = new AtomicReference<>(null); + private final AtomicLong taskCurrentRefreshIntervalMs; + + // Thundering herd guard + private final AtomicBoolean isRefreshing = new AtomicBoolean(false); + + // Circuit breaker + private final AtomicInteger consecutiveFailures = new AtomicInteger(0); + private final AtomicInteger circuitBreakerSkipCount = new AtomicInteger(0); + + QueueMetricsState(QueueMetricsSnapshot snapshot, long refreshIntervalMs) { + this.snapshot = new AtomicReference<>(snapshot); + this.lastWriteTime = 0L; // epoch = "never written" — ensures first fetch fires immediately + this.minRefreshIntervalMs = new AtomicLong(refreshIntervalMs); + this.taskCurrentRefreshIntervalMs = new AtomicLong(refreshIntervalMs); + } + + /** + * Returns the latest snapshot, or null if not yet fetched. + */ + public QueueMetricsSnapshot getSnapshot() { + return snapshot.get(); + } + + /** + * Returns ms since last successful RM write. Large value on first call (lastWriteTime=0). + */ + public long getAgeMs() { + return System.currentTimeMillis() - lastWriteTime; + } + + /** + * Returns the minimum refresh interval across all active sessions. + */ + public long getMinRefreshIntervalMs() { + return minRefreshIntervalMs.get(); + } + + /** Updates snapshot and lastWriteTime after a successful RM fetch. */ + public void applySnapshot(QueueMetricsSnapshot newSnapshot, + long refreshIntervalMs) { + this.snapshot.set(newSnapshot); + this.lastWriteTime = System.currentTimeMillis(); + minRefreshIntervalMs.updateAndGet(current -> Math.min(current, refreshIntervalMs)); + } + + + /** + * Registers this session's interval. Returns true if rescheduling may be needed + * (no task running, or this session lowered the minimum interval). + * Thread-safe: compute() is atomic per-key; getAndAccumulate returns previous value + * so only the thread that actually lowered the minimum triggers rescheduling. + */ + public boolean registerInterval(long refreshIntervalMs) { + intervalCounts.compute(refreshIntervalMs, (k, existing) -> { + if (existing == null) { + return new AtomicInteger(1); + } + existing.incrementAndGet(); + return existing; + }); + long prevMin = minRefreshIntervalMs.getAndAccumulate(refreshIntervalMs, Math::min); + int count = activeSessionCount.incrementAndGet(); + LOG.debug("Session registered at {}ms, activeCount={}", refreshIntervalMs, count); + return refreshTask.get() == null || refreshIntervalMs < prevMin; + } + + /** + * Deregisters this session's interval. Returns true if the task interval may need + * to change (this thread removed the last session at or below the current task interval). + * Thread-safe: compute() atomically decrements and conditionally removes the bucket. + */ + public boolean deregisterInterval(long refreshIntervalMs) { + boolean[] thisBucketRemoved = {false}; Review Comment: this is done by removing deregisterInterval at all, resolving this ########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/QueueMetricsRefreshPool.java: ########## @@ -0,0 +1,183 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; + +/** + * Singleton manager for the JVM-wide refresh executor pool used by queue metrics collection. + * Provides a shared {@link ScheduledExecutorService} that fires periodic YARN RM refresh tasks + * across all queries in the HiveServer2 process. + * <p> + * Initialized during HiveServer2 startup via {@link #init(int)} when Tez session pool is set up. + * Pool size is configured via {@code hive.server2.tez.queue.metrics.refresh.threads} (default: 4). + * <p> + * All {@link YarnQueueMetricsCollector} instances share this single pool, ensuring efficient + * resource usage and preventing thread explosion when many queries run concurrently. + * <p> + * Thread-safe singleton implementation using double-check locking pattern. + */ +public final class QueueMetricsRefreshPool { + private static final Logger LOG = LoggerFactory.getLogger(QueueMetricsRefreshPool.class); + + private static final int DEFAULT_THREAD_COUNT = 4; + public static final int JITTER_PERCENT = 10; + + private static final AtomicReference<QueueMetricsRefreshPool> instance = new AtomicReference<>(null); Review Comment: how do you mean "better for concurrent HiveServer2 startup"? when a HS2 is started, we don't expect a concurrent HS2 in the same JVM ########## service/src/test/org/apache/hive/service/server/TestHiveServer2QueueMetricsPoolInit.java: ########## @@ -0,0 +1,247 @@ +/* Review Comment: see my comment above, this doesn't have to be tested for "tez" ########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/YarnQueueMetricsCollector.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Objects; + +/** + * Collects YARN queue resource metrics using a shared cache to reduce ResourceManager load. + * Coordinates with other collectors via QueueMetricsCache to prevent duplicate RM calls. + * + * Executor pool management (sizing, lifecycle, JMX) is fully delegated to + * {@link QueueMetricsRefreshPool}. This class focuses solely on per-query/per-queue + * metrics logic: session registration, refresh scheduling, and cache coordination. + */ +public class YarnQueueMetricsCollector implements QueueMetricsCollector { + private static final Logger LOG = LoggerFactory.getLogger(YarnQueueMetricsCollector.class); + + private final YarnClient yarnClient; + private final String queueName; + private final long refreshIntervalMs; + private final String dagName; + + /** + * Creates a collector for the given queue and query. Non-blocking: metrics are + * fetched asynchronously; the first progress update may show no metrics, subsequent + * updates will once the first fetch completes (within one refreshIntervalMs). + * + * @param yarnClient Live YarnClient from the Tez session + * @param queueName YARN queue this query runs on + * @param refreshIntervalMs How often to poll YARN RM (ms) + * @param dagName Hive query identifier used for logging Review Comment: I think this is rather something like: "Tez DAG name used to identify the query for logging purposes." ########## ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/QueueMetricsRefreshPool.java: ########## @@ -0,0 +1,188 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Singleton manager for the JVM-wide refresh executor pool used by queue metrics collection. + * Provides a shared {@link ScheduledExecutorService} that fires periodic YARN RM refresh tasks + * across all queries in the HiveServer2 process. + * <p> + * Initialized during HiveServer2 startup via {@link #init(int)} when Tez session pool is set up. + * Pool size is configured via {@code hive.server2.tez.queue.metrics.refresh.threads} (default: 4). + * <p> + * All {@link YarnQueueMetricsCollector} instances share this single pool, ensuring efficient + * resource usage and preventing thread explosion when many queries run concurrently. + * <p> + * Thread-safe singleton implementation using double-check locking pattern. + */ +public final class QueueMetricsRefreshPool { + private static final Logger LOG = LoggerFactory.getLogger(QueueMetricsRefreshPool.class); + + private static final int DEFAULT_THREAD_COUNT = 4; + public static final int JITTER_PERCENT = 10; + + private static final AtomicReference<QueueMetricsRefreshPool> INSTANCE = new AtomicReference<>(null); Review Comment: in the latest code this is still an `AtomicReference` ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/TestYarnQueueMetricsCollector.java: ########## @@ -0,0 +1,587 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.api.records.QueueStatistics; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mockingDetails; +import static org.mockito.Mockito.when; + +/** + * Test cases for YarnQueueMetricsCollector. + */ +public class TestYarnQueueMetricsCollector { + + @Mock + private YarnClient mockYarnClient; + + @Mock + private QueueInfo mockQueueInfo; + + @Mock + private QueueStatistics mockQueueStats; + + private AutoCloseable closeable; + private HiveConf testConf; + + private static final long WAIT_TIMEOUT_MS = 5000; + + @Before + public void setUp() { + closeable = MockitoAnnotations.openMocks(this); + testConf = new HiveConf(); + // Reset the pool manager singleton and cache so each test starts with a clean state. + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + @After + public void tearDown() throws Exception { + if (closeable != null) { + closeable.close(); + } + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + /** + * Helper to create a collector in tests using a default HiveConf (min pool sizes). + */ + private YarnQueueMetricsCollector newCollector(YarnClient yarnClient, String queueName, + long refreshIntervalMs, String queryId) { + return new YarnQueueMetricsCollector(yarnClient, queueName, refreshIntervalMs, queryId, testConf); + } + + /** + * Waits for a snapshot to be available (non-null). + */ + private QueueMetricsSnapshot waitForSnapshot( + YarnQueueMetricsCollector collector, long timeoutMs) { + long startTime = System.currentTimeMillis(); Review Comment: see another comment regarding Awaitility ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/TestTezJobMonitorQueueMetrics.java: ########## @@ -0,0 +1,293 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring; + +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.NoOpQueueMetricsCollector; +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.QueueMetricsCollector; +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.YarnQueueMetricsCollector; + +import java.lang.reflect.Field; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConfForTest; +import org.apache.hadoop.hive.ql.Context; +import org.apache.hadoop.hive.ql.exec.tez.TezSession; +import org.apache.hadoop.hive.ql.log.PerfLogger; +import org.apache.hadoop.hive.ql.plan.BaseWork; +import org.apache.hadoop.hive.ql.session.SessionState; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.apache.tez.common.counters.TezCounters; +import org.apache.tez.dag.api.DAG; +import org.apache.tez.dag.api.client.DAGClient; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Test cases for TezJobMonitor queue metrics initialization. + */ +public class TestTezJobMonitorQueueMetrics { + + @Mock + private TezSession mockSession; + @Mock + private DAGClient mockDagClient; + @Mock + private DAG mockDag; + @Mock + private Context mockContext; + @Mock + private PerfLogger mockPerfLogger; + @Mock + private YarnClient mockYarnClient; + @Mock + private TezCounters mockCounters; + + private HiveConf hiveConf; + private List<BaseWork> topSortedWorks; + private SessionState sessionState; + private AutoCloseable mockCloseable; + + @Before + public void setUp() { + mockCloseable = MockitoAnnotations.openMocks(this); + hiveConf = new HiveConfForTest(TestTezJobMonitorQueueMetrics.class); + hiveConf.set("hive.security.authorization.manager", + "org.apache.hadoop.hive.ql.security.authorization.plugin.sqlstd.SQLStdConfOnlyAuthorizerFactory"); + sessionState = SessionState.start(hiveConf); + topSortedWorks = new ArrayList<>(); + when(mockDag.getName()).thenReturn("test-dag-1"); + } + + @After + public void tearDown() throws Exception { + if (mockCloseable != null) { + mockCloseable.close(); + } + if (sessionState != null) { + sessionState.close(); + } + } + + @Test + public void testMetricsCollectorDisabledByDefault() throws Exception { + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval=0), getYarnClient() is never called because + // the check happens before attempting to retrieve the YarnClient + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorEnabledWithInterval() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + verify(mockSession, atLeastOnce()).getYarnClient(); + } + + @Test + public void testMetricsCollectorDisabledWithZeroInterval() throws Exception { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 0, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval=0), getYarnClient() is never called + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorDisabledWithNegativeInterval() throws Exception { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, -1, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval<0), getYarnClient() is never called + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorWithSmallInterval() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 500, TimeUnit.MILLISECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created with adjusted interval", monitor); + } + + @Test + public void testMetricsCollectorWithCustomQueue() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 15, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("production.analytics"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + verify(mockSession, atLeastOnce()).getQueueName(); + assertNotNull("Monitor should be created with custom queue", monitor); + } + + /** + * Metrics enabled with a null YarnClient: monitor must still be created and must + * reach the YarnClient gate (verify getYarnClient called), but must NOT call + * getQueueName (nothing to resolve without a client). + */ + @Test + public void testMetricsCollectorWithNullYarnClient() { Review Comment: okay, this makes sense to me now after looking at TezSessionState ``` LOG.warn("Failed to initialize YarnClient for metrics collection", e); yarnClient = null; ``` given this is just a resource monitoring, and not a critical infrastructure for the actual query, falling back silently is okay, HS2 logs will contain the rest of the details ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezSessionState.java: ########## @@ -133,4 +137,122 @@ void openInternalUnsafe(boolean isAsync, SessionState.LogHelper console) { sessionStateForTest.open(resources); } + + /** + * Tests that YarnClient is NOT initialized when queue metrics are disabled (default: interval=0). + * This ensures zero overhead when the feature is disabled. + */ + @Test + public void testYarnClientNotInitializedWhenMetricsDisabled() { + SessionState ss = createSessionState(); + HiveConf hiveConf = ss.getConf(); + + // Default config: queue metrics disabled (interval = 0) + Assert.assertEquals("Default interval should be 0 (disabled)", + 0, HiveConf.getTimeVar(hiveConf, HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, TimeUnit.MILLISECONDS)); + + TezSessionState sessionState = new TezSessionState(ss.getSessionId(), hiveConf); + + // Mock a TezClient and set it + TezClient mockTezClient = Mockito.mock(TezClient.class); + sessionState.setTezClient(mockTezClient); + + // getYarnClient() should return null when metrics disabled + YarnClient yarnClient = sessionState.getYarnClient(); + Assert.assertNull("YarnClient should not be initialized when queue metrics are disabled", yarnClient); + } + + /** + * Tests that YarnClient IS lazily initialized when queue metrics are enabled. + * This ensures the client is created only when needed. + */ + @Test + public void testYarnClientLazilyInitializedWhenMetricsEnabled() { + SessionState ss = createSessionState(); + HiveConf hiveConf = ss.getConf(); + + // Enable queue metrics with a positive interval + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + + TezSessionState sessionState = new TezSessionState(ss.getSessionId(), hiveConf); + + // Mock a TezClient and set it + TezClient mockTezClient = Mockito.mock(TezClient.class); + sessionState.setTezClient(mockTezClient); + + // First call to getYarnClient() should initialize it + YarnClient yarnClient = sessionState.getYarnClient(); + Assert.assertNotNull("YarnClient should be initialized when queue metrics are enabled", yarnClient); + + // Second call should return the same instance + YarnClient yarnClient2 = sessionState.getYarnClient(); + Assert.assertSame("Should return the same YarnClient instance", yarnClient, yarnClient2); + } + + /** + * Tests that YarnClient is not initialized when TezClient is null, + * even if queue metrics are enabled. + */ + @Test + public void testYarnClientNotInitializedWhenTezClientNull() { + SessionState ss = createSessionState(); + HiveConf hiveConf = ss.getConf(); + + // Enable queue metrics + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + + TezSessionState sessionState = new TezSessionState(ss.getSessionId(), hiveConf); + + // Don't set TezClient (session is null) + + // getYarnClient() should return null when TezClient is not set + YarnClient yarnClient = sessionState.getYarnClient(); + Assert.assertNull("YarnClient should not be initialized when TezClient is null", yarnClient); + } + + /** + * Tests the thread-safety of lazy YarnClient initialization with concurrent calls. + */ + @Test + public void testYarnClientLazyInitializationThreadSafety() throws InterruptedException { + SessionState ss = createSessionState(); + HiveConf hiveConf = ss.getConf(); + + // Enable queue metrics + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + + TezSessionState sessionState = new TezSessionState(ss.getSessionId(), hiveConf); + TezClient mockTezClient = Mockito.mock(TezClient.class); + sessionState.setTezClient(mockTezClient); + + // Create multiple threads that call getYarnClient() concurrently + final int threadCount = 10; + Thread[] threads = new Thread[threadCount]; + YarnClient[] clients = new YarnClient[threadCount]; + + for (int i = 0; i < threadCount; i++) { + final int index = i; + threads[i] = new Thread(() -> { + clients[index] = sessionState.getYarnClient(); + }); + } + + // Start all threads + for (Thread thread : threads) { + thread.start(); + } + + // Wait for all threads to complete + for (Thread thread : threads) { + thread.join(); + } + + // All threads should get the same YarnClient instance + YarnClient firstClient = clients[0]; + Assert.assertNotNull("YarnClient should be initialized", firstClient); + + for (int i = 1; i < threadCount; i++) { + Assert.assertSame("All threads should get the same YarnClient instance", firstClient, clients[i]); + } + } Review Comment: thanks! nit: `getFirst` instead of `get(0)` (I know I advised that, it's just a warning now in my IDE :D ) ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/TestQueueMetricsState.java: ########## @@ -0,0 +1,322 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.api.records.QueueStatistics; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.when; + +/** + * Unit tests for QueueMetricsState - tests state management logic in isolation. + * Tests interval registration, circuit breaker, refresh locking, and other state logic. + */ +public class TestQueueMetricsState { + + @Mock + private QueueInfo mockQueueInfo; + + @Mock + private QueueStatistics mockQueueStats; + + + @Before + public void setUp() { + MockitoAnnotations.openMocks(this); + setupMockQueueInfo(); + } + + private void setupMockQueueInfo() { + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(1024L); + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(1024L); + when(mockQueueStats.getAllocatedVCores()).thenReturn(4L); + when(mockQueueStats.getAvailableVCores()).thenReturn(4L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(1L); + when(mockQueueStats.getNumAppsPending()).thenReturn(0L); + when(mockQueueStats.getAllocatedContainers()).thenReturn(2L); + when(mockQueueStats.getPendingContainers()).thenReturn(0L); + when(mockQueueInfo.getQueueStatistics()).thenReturn(mockQueueStats); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.25f); + } + + @Test + public void testConstructorWithNullSnapshot() { + QueueMetricsState state = new QueueMetricsState(null, 5000L); + + assertNull("Snapshot should be null when constructed with null", state.getSnapshot()); + assertEquals("Min interval should be set", 5000L, state.getMinRefreshIntervalMs()); + } + + @Test + public void testConstructorWithSnapshot() { + QueueMetricsSnapshot snapshot = new QueueMetricsSnapshot(mockQueueInfo); + QueueMetricsState state = new QueueMetricsState(snapshot, 10000L); + + assertNotNull("Snapshot should not be null", state.getSnapshot()); + assertEquals("Min interval should be set", 10000L, state.getMinRefreshIntervalMs()); + } + + @Test + public void testGetAgeMsReturnsLargeValueInitially() { + QueueMetricsState state = new QueueMetricsState(null, 5000L); + + long age = state.getAgeMs(); + + // Age should be very large when lastWriteTime = 0 (epoch) + assertTrue("Age should be > 1 year in ms", age > 365L * 24 * 60 * 60 * 1000); + } + + @Test + public void testApplySnapshotUpdatesSnapshot() { + QueueMetricsState state = new QueueMetricsState(null, 10000L); + assertNull("Initial snapshot should be null", state.getSnapshot()); + + QueueMetricsSnapshot snapshot = new QueueMetricsSnapshot(mockQueueInfo); + state.applySnapshot(snapshot, 5000L); + + assertNotNull("Snapshot should be updated", state.getSnapshot()); + assertEquals("Memory should match", 1.0f, state.getSnapshot().getMemoryUsedGB(), 0.01f); + } + + @Test + public void testApplySnapshotReducesAgeMs() { + QueueMetricsState state = new QueueMetricsState(null, 5000L); + long initialAge = state.getAgeMs(); + + // Spin-wait up to 200ms to ensure time has passed so the age comparison is meaningful + long deadline = System.currentTimeMillis() + 200; Review Comment: see another comment regarding Awaitility ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/TestTezJobMonitorQueueMetrics.java: ########## @@ -0,0 +1,293 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring; + +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.NoOpQueueMetricsCollector; +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.QueueMetricsCollector; +import org.apache.hadoop.hive.ql.exec.tez.monitoring.yarnqueue.YarnQueueMetricsCollector; + +import java.lang.reflect.Field; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConfForTest; +import org.apache.hadoop.hive.ql.Context; +import org.apache.hadoop.hive.ql.exec.tez.TezSession; +import org.apache.hadoop.hive.ql.log.PerfLogger; +import org.apache.hadoop.hive.ql.plan.BaseWork; +import org.apache.hadoop.hive.ql.session.SessionState; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.apache.tez.common.counters.TezCounters; +import org.apache.tez.dag.api.DAG; +import org.apache.tez.dag.api.client.DAGClient; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Test cases for TezJobMonitor queue metrics initialization. + */ +public class TestTezJobMonitorQueueMetrics { + + @Mock + private TezSession mockSession; + @Mock + private DAGClient mockDagClient; + @Mock + private DAG mockDag; + @Mock + private Context mockContext; + @Mock + private PerfLogger mockPerfLogger; + @Mock + private YarnClient mockYarnClient; + @Mock + private TezCounters mockCounters; + + private HiveConf hiveConf; + private List<BaseWork> topSortedWorks; + private SessionState sessionState; + private AutoCloseable mockCloseable; + + @Before + public void setUp() { + mockCloseable = MockitoAnnotations.openMocks(this); + hiveConf = new HiveConfForTest(TestTezJobMonitorQueueMetrics.class); + hiveConf.set("hive.security.authorization.manager", + "org.apache.hadoop.hive.ql.security.authorization.plugin.sqlstd.SQLStdConfOnlyAuthorizerFactory"); + sessionState = SessionState.start(hiveConf); + topSortedWorks = new ArrayList<>(); + when(mockDag.getName()).thenReturn("test-dag-1"); + } + + @After + public void tearDown() throws Exception { + if (mockCloseable != null) { + mockCloseable.close(); + } + if (sessionState != null) { + sessionState.close(); + } + } + + @Test + public void testMetricsCollectorDisabledByDefault() throws Exception { + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval=0), getYarnClient() is never called because + // the check happens before attempting to retrieve the YarnClient + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorEnabledWithInterval() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + verify(mockSession, atLeastOnce()).getYarnClient(); + } + + @Test + public void testMetricsCollectorDisabledWithZeroInterval() throws Exception { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 0, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval=0), getYarnClient() is never called + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorDisabledWithNegativeInterval() throws Exception { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, -1, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(null); // YarnClient should not be initialized + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created", monitor); + // When metrics are disabled (interval<0), getYarnClient() is never called + verify(mockSession, never()).getYarnClient(); + verify(mockYarnClient, never()).getQueueInfo(anyString()); + } + + @Test + public void testMetricsCollectorWithSmallInterval() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 500, TimeUnit.MILLISECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor should be created with adjusted interval", monitor); + } + + @Test + public void testMetricsCollectorWithCustomQueue() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 15, TimeUnit.SECONDS); + + when(mockSession.getYarnClient()).thenReturn(mockYarnClient); + when(mockSession.getQueueName()).thenReturn("production.analytics"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + verify(mockSession, atLeastOnce()).getQueueName(); + assertNotNull("Monitor should be created with custom queue", monitor); + } + + /** + * Metrics enabled with a null YarnClient: monitor must still be created and must + * reach the YarnClient gate (verify getYarnClient called), but must NOT call + * getQueueName (nothing to resolve without a client). + */ + @Test + public void testMetricsCollectorWithNullYarnClient() { + hiveConf.setTimeVar(HiveConf.ConfVars.HIVE_TEZ_QUEUE_METRICS_REFRESH_INTERVAL, 10, TimeUnit.SECONDS); + when(mockSession.getYarnClient()).thenReturn(null); + when(mockSession.getQueueName()).thenReturn("default"); + + TezJobMonitor monitor = + new TezJobMonitor(mockSession, topSortedWorks, mockDagClient, hiveConf, mockDag, mockContext, mockCounters, + mockPerfLogger); + + assertNotNull("Monitor must be created when YarnClient is null", monitor); + verify(mockSession, atLeastOnce()).getYarnClient(); + } + + /** + * Metrics enabled with a null queue name: monitor must be created and the code + * must reach both the YarnClient and queue-name gates. + */ + @Test + public void testMetricsCollectorWithNullQueueName() { Review Comment: I cannot sett `testMetricsCollectorWithNullQueueName` anymore, is it intentional? ########## service/src/java/org/apache/hive/service/server/HiveServer2.java: ########## @@ -966,6 +971,66 @@ private void initAndStartWorkloadManager(final WMFullResourcePlan resourcePlan) } } + /** + * Initializes the shared JVM-wide queue metrics refresh pool. + * <p> + * This pool provides background threads for periodic YARN queue metrics collection across + * all Tez sessions. The pool is shared by all queries in this HiveServer2 process to prevent + * thread explosion when many queries run concurrently. + * <p> + * Thread count is configured via {@code hive.server2.tez.queue.metrics.refresh.threads}. + * <p> + * The pool is only initialized when execution engine is "tez". Whether to actually collect + * queue metrics is controlled per-session by {@code hive.tez.queue.metrics.refresh.interval}, + * which is checked when creating metrics collectors for each query. + * <p> + * In non-Tez environments (MR, Spark, local), the pool is not created, avoiding unnecessary Review Comment: this is a bit confusing: "MR, Spark, local": the engine could not be "local", so it can be removed, also, "Spark" has already removed from master, only MR remains here, but I feel, this check and comment causes more noise than benefit, I mean: we don't have to guard accidental `MR` setups for performance: this whole "isTez" logic can be removed with the comments ########## ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/TestYarnQueueMetricsCollector.java: ########## @@ -0,0 +1,587 @@ +/* + * 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.hadoop.hive.ql.exec.tez.monitoring.yarnqueue; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.yarn.api.records.QueueInfo; +import org.apache.hadoop.yarn.api.records.QueueStatistics; +import org.apache.hadoop.yarn.client.api.YarnClient; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mockingDetails; +import static org.mockito.Mockito.when; + +/** + * Test cases for YarnQueueMetricsCollector. + */ +public class TestYarnQueueMetricsCollector { + + @Mock + private YarnClient mockYarnClient; + + @Mock + private QueueInfo mockQueueInfo; + + @Mock + private QueueStatistics mockQueueStats; + + private AutoCloseable closeable; + private HiveConf testConf; + + private static final long WAIT_TIMEOUT_MS = 5000; + + @Before + public void setUp() { + closeable = MockitoAnnotations.openMocks(this); + testConf = new HiveConf(); + // Reset the pool manager singleton and cache so each test starts with a clean state. + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + @After + public void tearDown() throws Exception { + if (closeable != null) { + closeable.close(); + } + QueueMetricsRefreshPool.resetForTesting(); + QueueMetricsCache.resetForTesting(); + } + + /** + * Helper to create a collector in tests using a default HiveConf (min pool sizes). + */ + private YarnQueueMetricsCollector newCollector(YarnClient yarnClient, String queueName, + long refreshIntervalMs, String queryId) { + return new YarnQueueMetricsCollector(yarnClient, queueName, refreshIntervalMs, queryId, testConf); + } + + /** + * Waits for a snapshot to be available (non-null). + */ + private QueueMetricsSnapshot waitForSnapshot( + YarnQueueMetricsCollector collector, long timeoutMs) { + long startTime = System.currentTimeMillis(); + QueueMetricsSnapshot snapshot; + while ((snapshot = collector.getLatestSnapshot()) == null) { + if (System.currentTimeMillis() - startTime > timeoutMs) { + fail("Snapshot not available after " + timeoutMs + "ms"); + } + Thread.onSpinWait(); // Hint to JVM that this is a spin-wait loop + } + return snapshot; + } + + /** + * Waits for a specific number of invocations with timeout. + */ + private void waitForInvocationCount(Object mock, int minCount, long timeoutMs) { + long startTime = System.currentTimeMillis(); + while (mockingDetails(mock).getInvocations().size() < minCount) { + if (System.currentTimeMillis() - startTime > timeoutMs) { + return; + } + Thread.onSpinWait(); // Hint to JVM that this is a spin-wait loop + } + } + + /** + * Helper method that configures mock objects with standard happy-path values. + */ + private void setupHappyPathMocks() throws Exception { + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(1024L); + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(1024L); + when(mockQueueStats.getAllocatedVCores()).thenReturn(4L); + when(mockQueueStats.getAvailableVCores()).thenReturn(4L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(1L); + when(mockQueueStats.getNumAppsPending()).thenReturn(0L); + when(mockQueueStats.getAllocatedContainers()).thenReturn(2L); + when(mockQueueStats.getPendingContainers()).thenReturn(0L); + when(mockQueueInfo.getQueueStatistics()).thenReturn(mockQueueStats); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.25f); + when(mockYarnClient.getQueueInfo(anyString())).thenReturn(mockQueueInfo); + } + + @Test(expected = IllegalArgumentException.class) + public void testConstructorWithNullYarnClient() { + new YarnQueueMetricsCollector(null, "default", 1000, "query-1", testConf); + } + + @Test(expected = IllegalArgumentException.class) + public void testConstructorWithNullQueueName() { + new YarnQueueMetricsCollector(mockYarnClient, null, 1000, "query-1", testConf); + } + + @Test + public void testSuccessfulMetricsCollection() throws Exception { + setupHappyPathMocks(); + when(mockYarnClient.getQueueInfo("default")).thenReturn(mockQueueInfo); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-1"); + try { + QueueMetricsSnapshot snapshot = waitForSnapshot(collector, WAIT_TIMEOUT_MS); + + assertNotNull("Snapshot should not be null", snapshot); + assertEquals("Memory used should be 1GB", 1.0f, snapshot.getMemoryUsedGB(), 0.1f); + assertEquals("Memory total should be 2GB (1+1)", 2.0f, snapshot.getMemoryTotalGB(), 0.1f); + assertEquals("VCores used should be 4", 4, snapshot.getVCoresUsed()); + assertEquals("VCores total should be 8 (4+4)", 8, snapshot.getVCoresTotal()); + assertEquals("Running apps should be 1", 1, snapshot.getRunningApps()); + assertEquals("Pending apps should be 0", 0, snapshot.getPendingApps()); + assertEquals("Allocated containers should be 2", 2, snapshot.getAllocatedContainers()); + assertEquals("Pending containers should be 0", 0, snapshot.getPendingContainers()); + assertEquals("Capacity should be 50%", 50.0f, snapshot.getCapacityPercentage(), 0.1f); + assertEquals("Current capacity should be 25%", 25.0f, snapshot.getCurrentCapacityPercentage(), 0.1f); + assertEquals("Memory percentage", "50.00%", snapshot.getMemoryPercentage()); + assertEquals("VCores percentage", "50.00%", snapshot.getVCoresPercentage()); + } finally { + collector.shutdown(); + } + } + + @Test + public void testMetricsCollectionWithNullQueueInfo() throws Exception { + when(mockYarnClient.getQueueInfo("nonexistent")).thenReturn(null); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "nonexistent", 10000, "test-query-2"); + try { + assertNull("Snapshot should be null for nonexistent queue", collector.getLatestSnapshot()); + } finally { + collector.shutdown(); + } + } + + @Test + public void testMetricsCollectionWithNullQueueStatistics() throws Exception { + when(mockQueueInfo.getQueueStatistics()).thenReturn(null); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.0f); + when(mockYarnClient.getQueueInfo("default")).thenReturn(mockQueueInfo); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-3"); + try { + QueueMetricsSnapshot snapshot = waitForSnapshot(collector, WAIT_TIMEOUT_MS); + assertNotNull("Snapshot should not be null", snapshot); + assertEquals("Memory used should be 0", 0.0f, snapshot.getMemoryUsedGB(), 0.01f); + assertEquals("Memory total should be 0", 0.0f, snapshot.getMemoryTotalGB(), 0.01f); + assertEquals("VCores used should be 0", 0, snapshot.getVCoresUsed()); + assertEquals("VCores total should be 0", 0, snapshot.getVCoresTotal()); + assertEquals("Capacity should still be 50%", 50.0f, snapshot.getCapacityPercentage(), 0.1f); + assertEquals("Current capacity should be 0%", 0.0f, snapshot.getCurrentCapacityPercentage(), 0.1f); + } finally { + collector.shutdown(); + } + } + + @Test + public void testPercentageCalculationWithZeroTotal() { + // Setup with zero totals + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(0L); + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(0L); + when(mockQueueStats.getAllocatedVCores()).thenReturn(0L); + when(mockQueueStats.getAvailableVCores()).thenReturn(0L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(0L); + when(mockQueueStats.getNumAppsPending()).thenReturn(0L); + when(mockQueueStats.getAllocatedContainers()).thenReturn(0L); + when(mockQueueStats.getPendingContainers()).thenReturn(0L); + when(mockQueueInfo.getQueueStatistics()).thenReturn(mockQueueStats); + when(mockQueueInfo.getCapacity()).thenReturn(0.0f); + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.0f); + + QueueMetricsSnapshot snapshot = + new QueueMetricsSnapshot(mockQueueInfo); + + // Should return "N/A" for percentages when total is zero + assertEquals("Memory percentage should be N/A", "N/A", snapshot.getMemoryPercentage()); + assertEquals("VCores percentage should be N/A", "N/A", snapshot.getVCoresPercentage()); + } + + @Test + public void testShutdownIdempotency() throws Exception { + when(mockYarnClient.getQueueInfo("default")).thenReturn(mockQueueInfo); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-4"); + collector.shutdown(); + collector.shutdown(); // second call must be safe + assertTrue("Multiple shutdowns should be safe", true); + } + + @Test + public void testExceptionDuringCollection() throws Exception { + when(mockYarnClient.getQueueInfo("default")) + .thenThrow(new RuntimeException("RM unavailable")); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-5"); + try { + assertNull("Snapshot should be null after exception", collector.getLatestSnapshot()); + } finally { + collector.shutdown(); + } + } + + @Test + public void testQueueNameRetrieval() throws Exception { + when(mockYarnClient.getQueueInfo(anyString())).thenReturn(mockQueueInfo); + when(mockQueueInfo.getQueueStatistics()).thenReturn(null); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "production", 10000, "test-query-6"); + try { + assertEquals("Queue name should match", "production", collector.getQueueName()); + } finally { + collector.shutdown(); + } + } + + @Test + public void testMemoryAndVCoreCalculation() { + // Test with specific values to verify calculation + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(5120L); // 5GB used + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(15360L); // 15GB available + when(mockQueueStats.getAllocatedVCores()).thenReturn(50L); + when(mockQueueStats.getAvailableVCores()).thenReturn(150L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(3L); + when(mockQueueStats.getNumAppsPending()).thenReturn(2L); + when(mockQueueStats.getAllocatedContainers()).thenReturn(10L); + when(mockQueueStats.getPendingContainers()).thenReturn(7L); + when(mockQueueInfo.getQueueStatistics()).thenReturn(mockQueueStats); + when(mockQueueInfo.getCapacity()).thenReturn(0.2f); // 20% + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.05f); // 5% + + QueueMetricsSnapshot snapshot = + new QueueMetricsSnapshot(mockQueueInfo); + + // Total = Used + Available + assertEquals("Memory used", 5.0f, snapshot.getMemoryUsedGB(), 0.01f); + assertEquals("Memory total", 20.0f, snapshot.getMemoryTotalGB(), 0.01f); // 5+15 + assertEquals("Memory percentage", "25.00%", snapshot.getMemoryPercentage()); // 5/20 + + assertEquals("VCores used", 50, snapshot.getVCoresUsed()); + assertEquals("VCores total", 200, snapshot.getVCoresTotal()); // 50+150 + assertEquals("VCores percentage", "25.00%", snapshot.getVCoresPercentage()); // 50/200 + + assertEquals("Running apps", 3, snapshot.getRunningApps()); + assertEquals("Pending apps", 2, snapshot.getPendingApps()); + assertEquals("Allocated containers", 10, snapshot.getAllocatedContainers()); + assertEquals("Pending containers", 7, snapshot.getPendingContainers()); + assertEquals("Capacity", 20.0f, snapshot.getCapacityPercentage(), 0.01f); + assertEquals("Current capacity", 5.0f, snapshot.getCurrentCapacityPercentage(), 0.01f); + } + + @Test(expected = IllegalArgumentException.class) + public void testQueueMetricsSnapshotWithNullQueueInfo() { + new QueueMetricsSnapshot(null); + } + + // ------------------------------------------------------------------------- + // Tests for Issue #1: Jitter on initial delay (Thundering Herd prevention) + // ------------------------------------------------------------------------- + // Note: Jitter is implicitly tested by all tests that successfully create collectors. + // Explicit jitter distribution testing would require reflection to access private + // scheduling details, which is fragile and not worth the maintenance cost. + + @Test + public void testExecutorCleanupOnInitializationFailure() throws Exception { + when(mockYarnClient.getQueueInfo(anyString())) + .thenThrow(new RuntimeException("Simulated RM failure during init")); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "init-fail-query"); + try { + assertNull("Snapshot should be null after init failure", collector.getLatestSnapshot()); + } finally { + collector.shutdown(); + } + } + + @Test + public void testCircuitBreakerActivatesAfterMaxFailures() throws Exception { + when(mockYarnClient.getQueueInfo(anyString())) + .thenThrow(new RuntimeException("YARN RM unavailable")); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 50, "circuit-breaker-query-1"); + try { + waitForInvocationCount(mockYarnClient, 6, 1000); + assertNull("Snapshot should be null when circuit breaker active", collector.getLatestSnapshot()); + int callCount = mockingDetails(mockYarnClient).getInvocations().size(); + assertTrue("Circuit breaker should reduce calls (got " + callCount + ")", callCount < 12); + } finally { + collector.shutdown(); + } + } + + @Test + public void testCircuitBreakerResetsOnSuccess() throws Exception { + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(4096L); + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(4096L); + when(mockQueueStats.getAllocatedVCores()).thenReturn(50L); + when(mockQueueStats.getAvailableVCores()).thenReturn(50L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(2L); + when(mockQueueStats.getNumAppsPending()).thenReturn(1L); + when(mockQueueStats.getAllocatedContainers()).thenReturn(5L); + when(mockQueueStats.getPendingContainers()).thenReturn(5L); + when(mockQueueInfo.getQueueStatistics()).thenReturn(mockQueueStats); + when(mockQueueInfo.getCapacity()).thenReturn(0.3f); + when(mockQueueInfo.getCurrentCapacity()).thenReturn(0.2f); + when(mockYarnClient.getQueueInfo(anyString())) + .thenThrow(new RuntimeException("Temporary RM failure")) + .thenThrow(new RuntimeException("Temporary RM failure")) + .thenThrow(new RuntimeException("Temporary RM failure")) + .thenThrow(new RuntimeException("Temporary RM failure")) + .thenThrow(new RuntimeException("Temporary RM failure")) + .thenReturn(mockQueueInfo); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 30, "circuit-breaker-recovery-query"); + try { + waitForInvocationCount(mockYarnClient, 3, 200); + assertNull("Snapshot should be null after circuit breaker activates", collector.getLatestSnapshot()); + QueueMetricsSnapshot snapshot = waitForSnapshot(collector, 2000); + assertNotNull("Snapshot should be populated after circuit breaker recovery", snapshot); + assertEquals("Memory used should be 4GB", 4.0f, snapshot.getMemoryUsedGB(), 0.1f); + } finally { + collector.shutdown(); + } + } + + + @Test + public void testNullQueueInfoDoesNotTriggerCircuitBreaker() throws Exception { + when(mockYarnClient.getQueueInfo(anyString())).thenReturn(null); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "nonexistent-queue", 50, "null-queueinfo-query"); + try { + waitForInvocationCount(mockYarnClient, 8, 800); + assertNull("Snapshot should remain null for null QueueInfo", collector.getLatestSnapshot()); + int callCount = mockingDetails(mockYarnClient).getInvocations().size(); + assertTrue("Null QueueInfo should NOT trigger circuit breaker (got " + callCount + " calls)", + callCount >= 8); + } finally { + collector.shutdown(); + } + } + + @Test + public void testSnapshotCollectionTimestampIsRecent() throws Exception { + setupHappyPathMocks(); + long beforeCreate = System.currentTimeMillis(); + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "timestamp-test"); + try { + QueueMetricsSnapshot snapshot = waitForSnapshot(collector, WAIT_TIMEOUT_MS); + long afterCollect = System.currentTimeMillis(); + assertNotNull("Snapshot should not be null", snapshot); + assertTrue("Timestamp should be >= creation time", snapshot.getCollectionTimestamp() >= beforeCreate); + assertTrue("Timestamp should be <= current time", snapshot.getCollectionTimestamp() <= afterCollect); + assertTrue("Timestamp should not be zero", snapshot.getCollectionTimestamp() > 0); + } finally { + collector.shutdown(); + } + } + + @Test + public void testRefreshIntervalRespected() throws Exception { + setupHappyPathMocks(); + when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(2048L); + when(mockQueueStats.getAvailableMemoryMB()).thenReturn(2048L); + when(mockQueueStats.getAllocatedVCores()).thenReturn(8L); + when(mockQueueStats.getAvailableVCores()).thenReturn(8L); + when(mockQueueStats.getNumAppsRunning()).thenReturn(2L); + when(mockQueueInfo.getCapacity()).thenReturn(0.6f); + + long intervalMs = 100; + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", intervalMs, "refresh-interval-test"); + try { + waitForSnapshot(collector, WAIT_TIMEOUT_MS); + int callsAfterFirst = mockingDetails(mockYarnClient).getInvocations().size(); + long toleranceMs = intervalMs + (long) (intervalMs * 0.2) + 300; + waitForInvocationCount(mockYarnClient, callsAfterFirst + 1, toleranceMs); + int callsAfterWait = mockingDetails(mockYarnClient).getInvocations().size(); + assertTrue("At least one refresh should have occurred within interval + tolerance", + callsAfterWait > callsAfterFirst); + } finally { + collector.shutdown(); + } + } + + @Test + public void testZeroRefreshIntervalIsRejected() throws Exception { + when(mockQueueInfo.getQueueStatistics()).thenReturn(null); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + when(mockYarnClient.getQueueInfo(anyString())).thenReturn(mockQueueInfo); + + assertThrows(IllegalArgumentException.class, () -> + new YarnQueueMetricsCollector(mockYarnClient, "default", 0, "zero-interval-test", testConf)); + } + + @Test + public void testNegativeRefreshIntervalIsRejected() throws Exception { + when(mockQueueInfo.getQueueStatistics()).thenReturn(null); + when(mockQueueInfo.getCapacity()).thenReturn(0.5f); + when(mockYarnClient.getQueueInfo(anyString())).thenReturn(mockQueueInfo); + + assertThrows(IllegalArgumentException.class, () -> + new YarnQueueMetricsCollector(mockYarnClient, "default", -1000, "negative-interval-test", testConf)); + } + + @Test + public void testJitterCalculationRange() { + long intervalMs = 2000; + long maxJitter = intervalMs * QueueMetricsRefreshPool.JITTER_PERCENT / 100; // 200ms + + // Test multiple queue names to ensure jitter is in range + String[] queues = {"default", "production", "batch", "analytics", "q" + "x".repeat(50)}; + for (String queueName : queues) { + long jitter = QueueMetricsRefreshPool.calculateJitter(queueName, intervalMs); + assertTrue("Jitter should be >= 0 for " + queueName, jitter >= 0); + assertTrue("Jitter should be < maxJitter (" + maxJitter + "ms) for " + queueName, + jitter < maxJitter); + } + } + + @Test + public void testJitterIsDeterministic() { + long intervalMs = 5000; + String queueName = "production-analytics"; + + long jitter1 = QueueMetricsRefreshPool.calculateJitter(queueName, intervalMs); + long jitter2 = QueueMetricsRefreshPool.calculateJitter(queueName, intervalMs); + long jitter3 = QueueMetricsRefreshPool.calculateJitter(queueName, intervalMs); + + assertEquals("Jitter should be deterministic (same queue → same jitter)", jitter1, jitter2); + assertEquals("Jitter should be deterministic across multiple calls", jitter2, jitter3); + } + + @Test + public void testMultipleSessionsShareCacheState() throws Exception { + setupHappyPathMocks(); + + // Create two collectors for the same queue + YarnQueueMetricsCollector collector1 = newCollector(mockYarnClient, "default", 5000, "query-1"); + YarnQueueMetricsCollector collector2 = newCollector(mockYarnClient, "default", 5000, "query-2"); + + try { + // Wait for first snapshot + QueueMetricsSnapshot snapshot1 = waitForSnapshot(collector1, WAIT_TIMEOUT_MS); + + // Second collector should get same snapshot from cache (not null) + QueueMetricsSnapshot snapshot2 = collector2.getLatestSnapshot(); + + assertNotNull("Second collector should get cached snapshot", snapshot2); + assertEquals("Both collectors should see same memory value", + snapshot1.getMemoryUsedGB(), snapshot2.getMemoryUsedGB(), 0.01f); + } finally { + collector1.shutdown(); + collector2.shutdown(); + } + } + + @Test + public void testDynamicReschedulingOnIntervalChange() throws Exception { + setupHappyPathMocks(); + + // Start with slow collector (10s) + YarnQueueMetricsCollector slowCollector = newCollector(mockYarnClient, "default", 10000, "slow-query"); + // Wait for first snapshot to confirm slow collector has stabilized + waitForSnapshot(slowCollector, WAIT_TIMEOUT_MS); + + // Add fast collector (1s) - should trigger rescheduling to 1s + YarnQueueMetricsCollector fastCollector = newCollector(mockYarnClient, "default", 1000, "fast-query"); + + try { + // Verify both collectors see updates (implies task running at faster interval) + QueueMetricsSnapshot snapshot = waitForSnapshot(fastCollector, WAIT_TIMEOUT_MS); + assertNotNull("Fast collector should get snapshot quickly", snapshot); + + // Shutdown fast collector - should reschedule back to slow interval + fastCollector.shutdown(); + // Wait up to 500ms for rescheduling to complete + waitForInvocationCount(mockYarnClient, mockingDetails(mockYarnClient).getInvocations().size(), 500); + + // Verify slow collector still works + assertNotNull("Slow collector should continue after fast shutdown", + slowCollector.getLatestSnapshot()); + } finally { + slowCollector.shutdown(); + } + } + + @Test + public void testCircuitBreakerProbeEvery10Ticks() throws Exception { + // Mock to always fail + when(mockYarnClient.getQueueInfo(anyString())) + .thenThrow(new RuntimeException("RM always failing")); + + YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 50, "probe-test"); + + try { + // Wait for circuit breaker to activate (5 failures) + waitForInvocationCount(mockYarnClient, 6, 1000); + int callsAfterActivation = mockingDetails(mockYarnClient).getInvocations().size(); + + // Wait for next ~12 ticks at 50ms interval — poll until invocation count stabilizes + waitForInvocationCount(mockYarnClient, callsAfterActivation + 2, 800); + int callsAfterWait = mockingDetails(mockYarnClient).getInvocations().size(); + + // Should have ~1 probe attempt in 10 ticks + int probeAttempts = callsAfterWait - callsAfterActivation; + assertTrue("Circuit breaker should allow ~1 probe per 10 ticks, got " + probeAttempts, + probeAttempts >= 0 && probeAttempts <= 2); Review Comment: 1. consider using larger timeouts in general: at some places I can see 800ms, which is so low, that eventually will lead to flaky tests in a slow precommit environment (having it a few seconds is totally fine) 2. what is it that the test actually measures? this looks very confusing: `probeAttempts >= 1 && probeAttempts <= 2` if it's like, "at least 1 probe must have occurred", then it should be only `probeAttempts > 0`, the "We allow up to 2 to account for timing variance" part just brings confusion here I believe -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
