AMBARI-15372. Zookeeper service check fails after removing a host (aonishuk)

ambari-server/src/main/java/org/apache/ambari/server/controller/AmbariManagementControllerImpl.java

@@ -518,7 +518,7 @@ public class AmbariManagementControllerImpl implements AmbariManagementControlle
       ServiceComponent sc = s.getServiceComponent(
           request.getComponentName());
 
-      setRestartRequiredServices(s, request.getHostname());
+      setRestartRequiredServices(s, request.getComponentName());
 
       Host host;
       try {
@@ -659,10 +659,12 @@ public class AmbariManagementControllerImpl implements AmbariManagementControlle
   }
 
   private void setRestartRequiredServices(
-          Service service, String hostName) throws AmbariException {
-
+          Service service, String componentName) throws AmbariException {
     Cluster cluster = service.getCluster();
     StackId stackId = cluster.getCurrentStackVersion();
+    if (service.getServiceComponent(componentName).isClientComponent()) {
+      return;
+    }
     Set<String> needRestartServices = ambariMetaInfo.getRestartRequiredServicesNames(
         stackId.getStackName(), stackId.getStackVersion());
 
@@ -670,16 +672,13 @@ public class AmbariManagementControllerImpl implements AmbariManagementControlle
       Map<String, ServiceComponent> m = service.getServiceComponents();
       for (Entry<String, ServiceComponent> entry : m.entrySet()) {
         ServiceComponent serviceComponent = entry.getValue();
-        if (serviceComponent.isMasterComponent()) {
-          Map<String, ServiceComponentHost> schMap = serviceComponent.getServiceComponentHosts();
-          //schMap will only contain hostName when deleting a host; when adding a host the test skips the loop
-          if(schMap.containsKey(hostName)) {
-            for (Entry<String, ServiceComponentHost> sch : schMap.entrySet()) {
-              ServiceComponentHost serviceComponentHost = sch.getValue();
-              serviceComponentHost.setRestartRequired(true);
-            }
+        Map<String, ServiceComponentHost> schMap = serviceComponent.getServiceComponentHosts();
+
+          for (Entry<String, ServiceComponentHost> sch : schMap.entrySet()) {
+            ServiceComponentHost serviceComponentHost = sch.getValue();
+            serviceComponentHost.setRestartRequired(true);
           }
-        }
+
       }
     }
   }
@@ -3066,7 +3065,7 @@ public class AmbariManagementControllerImpl implements AmbariManagementControlle
             "DISABLED/INIT/INSTALLED/INSTALL_FAILED/UNKNOWN state. Current=" + componentHost.getState() + ".");
       }
 
-      setRestartRequiredServices(service, request.getHostname());
+      setRestartRequiredServices(service, request.getComponentName());
 
       if (!safeToRemoveSCHs.containsKey(component)) {
         safeToRemoveSCHs.put(component, new HashSet<ServiceComponentHost>());

ambari-server/src/test/java/org/apache/ambari/server/controller/AmbariManagementControllerTest.java

@@ -9813,12 +9813,16 @@ public class AmbariManagementControllerTest {
     String datanode = "DATANODE";
     String hdfsClient = "HDFS_CLIENT";
     String zookeeperServer = "ZOOKEEPER_SERVER";
+    String zookeeperClient = "ZOOKEEPER_CLIENT";
+
     createServiceComponent(clusterName, hdfsService, namenode,
       State.INIT);
     createServiceComponent(clusterName, hdfsService, datanode,
       State.INIT);
     createServiceComponent(clusterName, zookeeperService, zookeeperServer,
       State.INIT);
+    createServiceComponent(clusterName, zookeeperService, zookeeperClient,
+      State.INIT);
 
     String host1 = "h1";
     String host2 = "h2";
@@ -9828,6 +9832,8 @@ public class AmbariManagementControllerTest {
     createServiceComponentHost(clusterName, hdfsService, datanode, host1, null);
     createServiceComponentHost(clusterName, zookeeperService, zookeeperServer, host1,
       null);
+    createServiceComponentHost(clusterName, zookeeperService, zookeeperClient, host1,
+      null);
 
     ServiceComponentHost zookeeperSch = null;
     for (ServiceComponentHost sch : cluster.getServiceComponentHosts(host1)) {
@@ -9838,11 +9844,16 @@ public class AmbariManagementControllerTest {
     assertFalse(zookeeperSch.isRestartRequired());
 
     addHostToCluster(host2, clusterName);
-    createServiceComponentHost(clusterName, zookeeperService, zookeeperServer, host2, null);
+    createServiceComponentHost(clusterName, zookeeperService, zookeeperClient, host2, null);
 
     assertFalse(zookeeperSch.isRestartRequired());  //No restart required if adding host
 
+    createServiceComponentHost(clusterName, zookeeperService, zookeeperServer, host2, null);
+
+    assertTrue(zookeeperSch.isRestartRequired());  //Add zk server required restart
+
     deleteServiceComponentHost(clusterName, zookeeperService, zookeeperServer, host2, null);
+    deleteServiceComponentHost(clusterName, zookeeperService, zookeeperClient, host2, null);
     deleteHost(host2);
 
     assertTrue(zookeeperSch.isRestartRequired());   //Restart if removing host!