AMBARI-16185. Express Ugrade stuck on getting status on 1600 node cluster (ncole)

ambari-server/src/main/java/org/apache/ambari/server/controller/internal/StageResourceProvider.java

@@ -30,7 +30,6 @@ import java.util.Set;
 import javax.inject.Inject;
 import javax.inject.Provider;
 
-import com.google.common.collect.Sets;
 import org.apache.ambari.server.StaticallyInject;
 import org.apache.ambari.server.actionmanager.HostRoleStatus;
 import org.apache.ambari.server.controller.AmbariManagementController;
@@ -45,15 +44,20 @@ import org.apache.ambari.server.controller.spi.Resource;
 import org.apache.ambari.server.controller.spi.ResourceAlreadyExistsException;
 import org.apache.ambari.server.controller.spi.SystemException;
 import org.apache.ambari.server.controller.spi.UnsupportedPropertyException;
+import org.apache.ambari.server.controller.utilities.PredicateHelper;
 import org.apache.ambari.server.orm.dao.HostRoleCommandDAO;
 import org.apache.ambari.server.orm.dao.HostRoleCommandStatusSummaryDTO;
 import org.apache.ambari.server.orm.dao.StageDAO;
 import org.apache.ambari.server.orm.entities.StageEntity;
 import org.apache.ambari.server.state.Cluster;
 import org.apache.ambari.server.state.Clusters;
+import org.apache.ambari.server.topology.LogicalRequest;
 import org.apache.ambari.server.topology.TopologyManager;
 import org.apache.ambari.server.utils.SecretReference;
 import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang.math.NumberUtils;
+
+import com.google.common.collect.Sets;
 
 /**
  * ResourceProvider for Stage
@@ -217,11 +221,31 @@ public class StageResourceProvider extends AbstractControllerResourceProvider im
 
     cache.clear();
 
-    Collection<StageEntity> topologyManagerStages = topologyManager.getStages();
-    for (StageEntity entity : topologyManagerStages) {
-      Resource stageResource = toResource(entity, propertyIds);
-      if (predicate.evaluate(stageResource)) {
-        results.add(stageResource);
+    // !!! check the id passed to see if it's a LogicalRequest.  This safeguards against
+    // iterating all stages for all requests.  That is a problem when the request
+    // is for an Upgrade, but was pulling the data anyway.
+    Map<String, Object> map = PredicateHelper.getProperties(predicate);
+
+    if (map.containsKey(STAGE_REQUEST_ID)) {
+      Long requestId = NumberUtils.toLong(map.get(STAGE_REQUEST_ID).toString());
+      LogicalRequest lr = topologyManager.getRequest(requestId);
+
+      if (null != lr) {
+        Collection<StageEntity> topologyManagerStages = lr.getStageEntities();
+        for (StageEntity entity : topologyManagerStages) {
+          Resource stageResource = toResource(entity, propertyIds);
+          if (predicate.evaluate(stageResource)) {
+            results.add(stageResource);
+          }
+        }
+      }
+    } else {
+      Collection<StageEntity> topologyManagerStages = topologyManager.getStages();
+      for (StageEntity entity : topologyManagerStages) {
+        Resource stageResource = toResource(entity, propertyIds);
+        if (predicate.evaluate(stageResource)) {
+          results.add(stageResource);
+        }
       }
     }
 

ambari-server/src/main/java/org/apache/ambari/server/topology/TopologyManager.java

@@ -411,8 +411,10 @@ public class TopologyManager {
     }
   }
 
-  // currently we are just returning all stages for all requests
-  //and relying on the StageResourceProvider to convert each to a resource and do a predicate eval on each
+  /**
+   * Currently we are just returning all stages for all requests
+   * and relying on the StageResourceProvider to convert each to a resource and do a predicate eval on each.
+   */
   public Collection<StageEntity> getStages() {
     ensureInitialized();
     Collection<StageEntity> stages = new ArrayList<StageEntity>();

ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StageResourceProviderTest.java

@@ -22,6 +22,7 @@ import static org.easymock.EasyMock.createNiceMock;
 import static org.easymock.EasyMock.createStrictMock;
 import static org.easymock.EasyMock.expect;
 import static org.easymock.EasyMock.replay;
+import static org.easymock.EasyMock.reset;
 import static org.easymock.EasyMock.verify;
 import static org.junit.Assert.fail;
 
@@ -84,7 +85,7 @@ public class StageResourceProviderTest {
     hrcDao = createStrictMock(HostRoleCommandDAO.class);
     topologyManager = EasyMock.createNiceMock(TopologyManager.class);
 
-    expect(topologyManager.getStages()).andReturn(new ArrayList<StageEntity>()).anyTimes();
+    expect(topologyManager.getStages()).andReturn(new ArrayList<StageEntity>()).atLeastOnce();
 
     expect(hrcDao.findAggregateCounts(EasyMock.anyObject(Long.class))).andReturn(
         new HashMap<Long, HostRoleCommandStatusSummaryDTO>() {
@@ -174,6 +175,39 @@ public class StageResourceProviderTest {
     verify(dao, clusters, cluster);
   }
 
+  @Test
+  public void testGetResourcesWithRequest() throws Exception {
+    StageResourceProvider provider = new StageResourceProvider(managementController);
+
+    Request request = createNiceMock(Request.class);
+    Predicate predicate = new PredicateBuilder().property(StageResourceProvider.STAGE_REQUEST_ID).equals(1L).toPredicate();
+
+    List<StageEntity> entities = getStageEntities(HostRoleStatus.COMPLETED);
+
+    expect(dao.findAll(request, predicate)).andReturn(entities);
+
+    expect(clusters.getClusterById(anyLong())).andReturn(cluster).anyTimes();
+    expect(cluster.getClusterName()).andReturn("c1").anyTimes();
+    reset(topologyManager);
+    expect(topologyManager.getRequest(EasyMock.anyLong())).andReturn(null).atLeastOnce();
+
+    replay(topologyManager, dao, clusters, cluster, request);
+
+    Set<Resource> resources = provider.getResources(request, predicate);
+
+    Assert.assertEquals(1, resources.size());
+
+    Resource resource = resources.iterator().next();
+
+    Assert.assertEquals(100.0, resource.getPropertyValue(StageResourceProvider.STAGE_PROGRESS_PERCENT));
+    Assert.assertEquals(HostRoleStatus.COMPLETED, resource.getPropertyValue(StageResourceProvider.STAGE_STATUS));
+    Assert.assertEquals(HostRoleStatus.COMPLETED, resource.getPropertyValue(StageResourceProvider.STAGE_DISPLAY_STATUS));
+    Assert.assertEquals(1000L, resource.getPropertyValue(StageResourceProvider.STAGE_START_TIME));
+    Assert.assertEquals(2500L, resource.getPropertyValue(StageResourceProvider.STAGE_END_TIME));
+
+    verify(topologyManager, dao, clusters, cluster);
+  }
+
   /**
    * Tests getting the display status of a stage which can differ from the final
    * status.