(tez) branch master updated: TEZ-4722: Issue with NODE_TOTAL_COUNT counter (#507) (Raghav Aggarwal reviewed by Laszlo Bodor)

[abstractdog]

This is an automated email from the ASF dual-hosted git repository.

abstractdog pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/tez.git


The following commit(s) were added to refs/heads/master by this push:
     new 81e0a248e TEZ-4722: Issue with NODE_TOTAL_COUNT counter (#507) (Raghav 
Aggarwal reviewed by Laszlo Bodor)
81e0a248e is described below

commit 81e0a248efc1c9f5a839be23ad8d2980bd6b9ba8
Author: Raghav Aggarwal <raghavaggarwal03...@gmail.com>
AuthorDate: Mon Jun 22 19:13:43 2026 +0530

    TEZ-4722: Issue with NODE_TOTAL_COUNT counter (#507) (Raghav Aggarwal 
reviewed by Laszlo Bodor)
---
 .../tez/dag/app/rm/TaskSchedulerManager.java       |  4 ++--
 .../tez/dag/app/rm/TestTaskSchedulerManager.java   | 25 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git 
a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java 
b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java
index 0acc16957..b35cb3c00 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerManager.java
@@ -227,8 +227,8 @@ public class TaskSchedulerManager extends AbstractService 
implements
     return getNumClusterNodes(false);
   }
 
-  public int getNumClusterNodes(boolean tryUpdate){
-    if (cachedNodeCount == -1 && tryUpdate){
+  public int getNumClusterNodes(boolean tryUpdate) {
+    if (cachedNodeCount == -1 || tryUpdate) {
       cachedNodeCount = countAllNodes();
     }
     return cachedNodeCount;
diff --git 
a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java 
b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java
index 595c045d8..d81426020 100644
--- 
a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java
+++ 
b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerManager.java
@@ -1240,4 +1240,29 @@ public class TestTaskSchedulerManager {
       throw new TimeoutException("Timed out waiting for condition.");
     }
   }
+
+  @Test(timeout = 5000)
+  public void testGetNumClusterNodes() throws Exception {
+    Configuration conf = new Configuration(false);
+    taskSchedulerManager.init(conf);
+    taskSchedulerManager.start();
+
+    // Mock the underlying task scheduler to simulate a scale up from 10 nodes 
to 20 nodes
+    
when(mockTaskScheduler.getClusterNodeCount()).thenReturn(10).thenReturn(20);
+
+    // Initial call, cachedNodeCount is -1, should fetch and cache 10
+    int count1 = taskSchedulerManager.getNumClusterNodes(false);
+    assertEquals(10, count1);
+
+    // Second call, tryUpdate is false, should return cached 10
+    int count2 = taskSchedulerManager.getNumClusterNodes(false);
+    assertEquals(10, count2);
+
+    // Third call, tryUpdate is true, should fetch new value 20
+    int count3 = taskSchedulerManager.getNumClusterNodes(true);
+    assertEquals(20, count3);
+
+    // Verify getClusterNodeCount was called exactly twice (once initially, 
once on forced update)
+    verify(mockTaskScheduler, times(2)).getClusterNodeCount();
+  }
 }