AMBARI-8573. Upgrade Execute: add service check tasks (ncole)

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

@@ -375,8 +375,8 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       for (StageWrapper wrapper : group.items) {
         UpgradeItemEntity itemEntity = new UpgradeItemEntity();
         itemEntity.setText(wrapper.getText());
-        itemEntity.setTasks(wrapper.getHostsJson());
-        itemEntity.setHosts(wrapper.getTasksJson());
+        itemEntity.setTasks(wrapper.getTasksJson());
+        itemEntity.setHosts(wrapper.getHostsJson());
         itemEntities.add(itemEntity);
 
         // upgrade items match a stage
@@ -420,10 +420,16 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
   private void createStage(Cluster cluster, RequestStageContainer request, final String version,
       UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
 
-    if (wrapper.hasCommand()) {
-      makeRestartStage(cluster, request, version, entity, wrapper);
-    } else {
-      makeActionStage(cluster, request, version, entity, wrapper);
+    switch (wrapper.getType()) {
+      case RESTART:
+        makeRestartStage(cluster, request, version, entity, wrapper);
+        break;
+      case RU_TASKS:
+        makeActionStage(cluster, request, version, entity, wrapper);
+        break;
+      case SERVICE_CHECK:
+        makeServiceCheckStage(cluster, request, version, entity, wrapper);
+        break;
     }
 
   }
@@ -431,6 +437,19 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
   private void makeActionStage(Cluster cluster, RequestStageContainer request, final String version,
       UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
 
+    // add each host to this stage
+    RequestResourceFilter filter = new RequestResourceFilter("", "",
+        new ArrayList<String>(wrapper.getHosts()));
+
+    Map<String, String> params = new HashMap<String, String>();
+    params.put("tasks", entity.getTasks());
+
+    ActionExecutionContext actionContext = new ActionExecutionContext(
+        cluster.getClusterName(), "ru_execute_tasks",
+        Collections.singletonList(filter),
+        params);
+    actionContext.setTimeout(Short.valueOf((short)60));
+
     Map<String, String> hostLevelParams = new HashMap<String, String>();
     hostLevelParams.put(JDK_LOCATION, getManagementController().getJdkResourceUrl());
 
@@ -449,20 +468,6 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     stage.setStageId(stageId);
     entity.setStageId(Long.valueOf(stageId));
 
-    // add each host to this stage
-    RequestResourceFilter filter = new RequestResourceFilter("", "",
-        new ArrayList<String>(wrapper.getHosts()));
-
-    // !!! TODO when the custom action is underway, change this
-    Map<String, String> params = new HashMap<String, String>();
-    params.put("tasks", entity.getTasks());
-
-    ActionExecutionContext actionContext = new ActionExecutionContext(
-        cluster.getClusterName(), "ru_execute_tasks",
-        Collections.singletonList(filter),
-        params);
-    actionContext.setTimeout(Short.valueOf((short)60));
-
     // !!! TODO verify the action is valid
 
     actionExecutionHelper.get().addExecutionCommandsToStage(actionContext, stage);
@@ -527,5 +532,48 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     request.addStages(Collections.singletonList(stage));
   }
 
+  private void makeServiceCheckStage(Cluster cluster, RequestStageContainer request, String version,
+      UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
+
+    List<RequestResourceFilter> filters = new ArrayList<RequestResourceFilter>();
+
+    for (TaskWrapper tw : wrapper.getTasks()) {
+      filters.add(new RequestResourceFilter(tw.getService(), "", Collections.<String>emptyList()));
+    }
+
+    Map<String, String> restartCommandParams = new HashMap<String, String>();
+    restartCommandParams.put("version", version);
+
+    ActionExecutionContext actionContext = new ActionExecutionContext(
+        cluster.getClusterName(), "SERVICE_CHECK",
+        filters,
+        restartCommandParams);
+    actionContext.setTimeout(Short.valueOf((short)-1));
+
+    ExecuteCommandJson jsons = commandExecutionHelper.get().getCommandJson(
+        actionContext, cluster);
+
+    Stage stage = stageFactory.get().createNew(request.getId().longValue(),
+        "/tmp/ambari",
+        cluster.getClusterName(),
+        cluster.getClusterId(),
+        entity.getText(),
+        jsons.getClusterHostInfo(),
+        jsons.getCommandParamsForStage(),
+        jsons.getHostParamsForStage());
+
+    long stageId = request.getLastStageId() + 1;
+    if (0L == stageId) {
+      stageId = 1L;
+    }
+    stage.setStageId(stageId);
+    entity.setStageId(Long.valueOf(stageId));
+
+    Map<String, String> requestParams = new HashMap<String, String>();
+
+    commandExecutionHelper.get().addExecutionCommandsToStage(actionContext, stage, requestParams);
+
+    request.addStages(Collections.singletonList(stage));
+  }
 
 }

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ColocatedGrouping.java

@@ -19,6 +19,7 @@ package org.apache.ambari.server.state.stack.upgrade;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -29,6 +30,7 @@ import javax.xml.bind.annotation.XmlElement;
 import javax.xml.bind.annotation.XmlType;
 
 import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -57,6 +59,7 @@ public class ColocatedGrouping extends Grouping {
     private Map<String, List<TaskProxy>> initialBatch = new LinkedHashMap<String, List<TaskProxy>>();
     private Map<String, List<TaskProxy>> finalBatches = new LinkedHashMap<String, List<TaskProxy>>();
 
+
     private MultiHomedHolder(Batch batch) {
       this.batch = batch;
     }
@@ -83,6 +86,7 @@ public class ColocatedGrouping extends Grouping {
           proxy = new TaskProxy();
           proxy.message = getStageText("Preparing", pc.name, Collections.singleton(host));
           proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), pc.preTasks));
+          proxy.service = service;
           proxy.component = pc.name;
           targetList.add(proxy);
         }
@@ -94,6 +98,7 @@ public class ColocatedGrouping extends Grouping {
             proxy = new TaskProxy();
             proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), t));
             proxy.restart = true;
+            proxy.service = service;
             proxy.component = pc.name;
             proxy.message = getStageText("Restarting ", pc.name, Collections.singleton(host));
 
@@ -104,6 +109,7 @@ public class ColocatedGrouping extends Grouping {
         if (null != pc.postTasks && pc.postTasks.size() > 0) {
           proxy = new TaskProxy();
           proxy.component = pc.name;
+          proxy.service = service;
           proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), pc.postTasks));
           proxy.message = getStageText("Completing", pc.name, Collections.singleton(host));
           targetList.add(proxy);
@@ -123,21 +129,24 @@ public class ColocatedGrouping extends Grouping {
 
       results.addAll(fromProxies(initialBatch));
 
+      // !!! TODO when manual tasks are ready
 //      StageWrapper wrapper = new StageWrapper(
 //      ManualTask task = new ManualTask();
 //      task.message = batch.message;
 //      wrapper.tasks.add(new TaskWrapper(null, null, null, task));
-      // !!! TODO
 //      results.add(wrapper);
 
       results.addAll(fromProxies(finalBatches));
 
+
       return results;
     }
 
     private List<StageWrapper> fromProxies(Map<String, List<TaskProxy>> wrappers) {
       List<StageWrapper> results = new ArrayList<StageWrapper>();
 
+      Set<String> serviceChecks = new HashSet<String>();
+
       for (Entry<String, List<TaskProxy>> entry : wrappers.entrySet()) {
 
         // !!! stage per host, per type
@@ -145,13 +154,15 @@ public class ColocatedGrouping extends Grouping {
         StageWrapper execwrapper = null;
 
         for (TaskProxy t : entry.getValue()) {
+          serviceChecks.add(t.service);
+
           if (!t.restart) {
             if (null == wrapper) {
-              wrapper = new StageWrapper(t.message, t.tasks);
+              wrapper = new StageWrapper(StageWrapper.Type.RU_TASKS, t.message, t.getTasksArray());
             }
           } else {
             if (null == execwrapper) {
-              execwrapper = new StageWrapper(t.message, t.tasks);
+              execwrapper = new StageWrapper(StageWrapper.Type.RESTART, t.message, t.getTasksArray());
             }
           }
         }
@@ -165,6 +176,21 @@ public class ColocatedGrouping extends Grouping {
         }
       }
 
+      if (serviceChecks.size() > 0) {
+        // !!! add the service check task
+        List<TaskWrapper> tasks = new ArrayList<TaskWrapper>();
+        for (String service : serviceChecks) {
+          tasks.add(new TaskWrapper(service, "", Collections.<String>emptySet(), new ServiceCheckTask()));
+        }
+
+        StageWrapper wrapper = new StageWrapper(
+            StageWrapper.Type.SERVICE_CHECK,
+            "Service Check " + StringUtils.join(serviceChecks, ", "),
+            tasks.toArray(new TaskWrapper[tasks.size()]));
+
+        results.add(wrapper);
+      }
+
       return results;
     }
 
@@ -175,6 +201,7 @@ public class ColocatedGrouping extends Grouping {
    */
   private static class TaskProxy {
     private boolean restart = false;
+    private String service;
     private String component;
     private String message;
     private List<TaskWrapper> tasks = new ArrayList<TaskWrapper>();
@@ -188,6 +215,10 @@ public class ColocatedGrouping extends Grouping {
 
       return s;
     }
+
+    private TaskWrapper[] getTasksArray() {
+      return tasks.toArray(new TaskWrapper[0]);
+    }
   }
 
 }

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Grouping.java

@@ -19,6 +19,7 @@ package org.apache.ambari.server.state.stack.upgrade;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 
@@ -28,6 +29,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.ambari.server.state.stack.UpgradePack;
 import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+import org.apache.commons.lang.StringUtils;
 
 /**
  *
@@ -55,13 +57,15 @@ public class Grouping {
   private static class DefaultBuilder extends StageWrapperBuilder {
 
     private List<StageWrapper> stages = new ArrayList<StageWrapper>();
+    private Set<String> serviceChecks = new HashSet<String>();
 
     @Override
     public void add(Set<String> hosts, String service, ProcessingComponent pc) {
       if (null != pc.preTasks && pc.preTasks.size() > 0) {
         StageWrapper stage = new StageWrapper(
+            StageWrapper.Type.RU_TASKS,
             getStageText("Preparing", pc.name, hosts),
-            Collections.singletonList(new TaskWrapper(service, pc.name, hosts, pc.preTasks)));
+            new TaskWrapper(service, pc.name, hosts, pc.preTasks));
         stages.add(stage);
       }
 
@@ -71,8 +75,9 @@ public class Grouping {
         if (RestartTask.class.isInstance(t)) {
           for (String hostName : hosts) {
             StageWrapper stage = new StageWrapper(
+                StageWrapper.Type.RESTART,
                 getStageText("Restarting", pc.name, Collections.singleton(hostName)),
-                Collections.singletonList(new TaskWrapper(service, pc.name, Collections.singleton(hostName), t)));
+                new TaskWrapper(service, pc.name, Collections.singleton(hostName), t));
             stages.add(stage);
           }
         }
@@ -80,14 +85,35 @@ public class Grouping {
 
       if (null != pc.postTasks && pc.postTasks.size() > 0) {
         StageWrapper stage = new StageWrapper(
+            StageWrapper.Type.RU_TASKS,
             getStageText("Completing", pc.name, hosts),
-            Collections.singletonList(new TaskWrapper(service, pc.name, hosts, pc.postTasks)));
+            new TaskWrapper(service, pc.name, hosts, pc.postTasks));
         stages.add(stage);
       }
+
+      serviceChecks.add(service);
+
     }
 
     @Override
     public List<StageWrapper> build() {
+
+      List<TaskWrapper> tasks = new ArrayList<TaskWrapper>();
+      for (String service : serviceChecks) {
+        tasks.add(new TaskWrapper(
+            service, "", Collections.<String>emptySet(), new ServiceCheckTask()));
+      }
+
+      if (serviceChecks.size() > 0) {
+        StageWrapper wrapper = new StageWrapper(
+            StageWrapper.Type.SERVICE_CHECK,
+            "Service Check " + StringUtils.join(serviceChecks, ", "),
+            tasks.toArray(new TaskWrapper[0])
+            );
+
+        stages.add(wrapper);
+      }
+
       return stages;
     }
 

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ServiceCheckTask.java

@@ -0,0 +1,41 @@
+/**
+ * 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.ambari.server.state.stack.upgrade;
+
+import javax.xml.bind.annotation.XmlAccessType;
+import javax.xml.bind.annotation.XmlAccessorType;
+import javax.xml.bind.annotation.XmlRootElement;
+import javax.xml.bind.annotation.XmlTransient;
+import javax.xml.bind.annotation.XmlType;
+
+/**
+ * Used to represent a restart of a component.
+ */
+@XmlRootElement
+@XmlAccessorType(XmlAccessType.FIELD)
+@XmlType(name="service-check")
+public class ServiceCheckTask extends Task {
+
+  @XmlTransient
+  private Task.Type type = Task.Type.SERVICE_CHECK;
+
+  @Override
+  public Task.Type getType() {
+    return type;
+  }
+}

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java

@@ -18,7 +18,7 @@
 package org.apache.ambari.server.state.stack.upgrade;
 
 import java.util.ArrayList;
-import java.util.Collection;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -32,10 +32,16 @@ public class StageWrapper {
 
   private static Gson gson = new Gson();
   private String text;
+  private Type type;
 
   private List<TaskWrapper> tasks;
 
-  public StageWrapper(String text, List<TaskWrapper> tasks) {
+  public StageWrapper(Type type, String text, TaskWrapper... tasks) {
+    this(type, text, Arrays.asList(tasks));
+  }
+
+  private StageWrapper(Type type, String text, List<TaskWrapper> tasks) {
+    this.type = type;
     this.text = text;
     this.tasks = tasks;
   }
@@ -59,23 +65,6 @@ public class StageWrapper {
     return gson.toJson(realTasks);
   }
 
-  /**
-   * @return {@code true} if any of the tasks is a command type.  This affects
-   * the type of stage that is created.
-   */
-  public boolean hasCommand() {
-
-    for (TaskWrapper tw : tasks) {
-      for (Task t : tw.getTasks()) {
-        if (t.getType().isCommand()) {
-          return true;
-        }
-      }
-    }
-
-    return false;
-  }
-
   /**
    * @return the set of hosts for all tasks
    */
@@ -102,4 +91,24 @@ public class StageWrapper {
     return text;
   }
 
+  /**
+   * Gets the type of stage.  All tasks defined for the stage execute this type.
+   * @return the type
+   */
+  public Type getType() {
+    return type;
+  }
+
+  /**
+   * Indicates the type of wrapper.
+   */
+  public enum Type {
+    MANUAL,
+    RESTART,
+    RU_TASKS,
+    SERVICE_CHECK
+  }
+
+
+
 }

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java

@@ -55,7 +55,11 @@ public abstract class Task {
     /**
      * Task that is a restart command.
      */
-    RESTART;
+    RESTART,
+    /**
+     * Task that is a service check
+     */
+    SERVICE_CHECK;
 
     /**
      * @return {@code true} if the task is manual or automated.

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

@@ -18,8 +18,8 @@
 package org.apache.ambari.server.controller.internal;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 
 import java.lang.reflect.Field;
 import java.util.Collections;
@@ -132,6 +132,8 @@ public class UpgradeResourceProviderTest {
     ServiceComponent component = service.addServiceComponent("ZOOKEEPER_SERVER");
     component.addServiceComponentHost("h1");
 
+    component = service.addServiceComponent("ZOOKEEPER_CLIENT");
+    component.addServiceComponentHost("h1");
 
   }
 
@@ -167,11 +169,12 @@ public class UpgradeResourceProviderTest {
 
     UpgradeGroupEntity group = entity.getUpgradeGroups().get(0);
 
-    assertEquals(3, group.getItems().size());
+    assertEquals(4, group.getItems().size());
 
     assertTrue(group.getItems().get(0).getText().contains("Preparing"));
     assertTrue(group.getItems().get(1).getText().contains("Restarting"));
     assertTrue(group.getItems().get(2).getText().contains("Completing"));
+    assertTrue(group.getItems().get(3).getText().contains("Service Check"));
 
     ActionManager am = injector.getInstance(ActionManager.class);
     List<Long> requests = am.getRequestsByStatus(RequestStatus.IN_PROGRESS, 100, true);
@@ -181,7 +184,7 @@ public class UpgradeResourceProviderTest {
 
 
     List<Stage> stages = am.getRequestStatus(requests.get(0).longValue());
-    assertEquals(3, stages.size());
+    assertEquals(4, stages.size());
     for (int i = 0; i < stages.size(); i++) {
       Stage stage = stages.get(i);
       UpgradeItemEntity upgradeItem = group.getItems().get(i);
@@ -189,9 +192,8 @@ public class UpgradeResourceProviderTest {
       assertEquals(UpgradeState.NONE, upgradeItem.getState());
     }
 
-
     List<HostRoleCommand> tasks = am.getRequestTasks(requests.get(0).longValue());
-    assertEquals(3, tasks.size());
+    assertEquals(4, tasks.size());
 
     return status;
   }
@@ -257,7 +259,7 @@ public class UpgradeResourceProviderTest {
     ResourceProvider upgradeItemResourceProvider = new UpgradeItemResourceProvider(amc);
     resources = upgradeItemResourceProvider.getResources(request, predicate);
 
-    assertEquals(3, resources.size());
+    assertEquals(4, resources.size());
     res = resources.iterator().next();
     assertNotNull(res.getPropertyValue("UpgradeItem/status"));
 

ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java

@@ -82,9 +82,9 @@ public class UpgradeHelperTest {
     assertEquals("CORE_MASTER", groups.get(1).name);
     assertEquals("CORE_SLAVES", groups.get(2).name);
 
-    assertEquals(5, groups.get(0).items.size());
-    assertEquals(4, groups.get(1).items.size());
-    assertEquals(4, groups.get(2).items.size());
+    assertEquals(6, groups.get(0).items.size());
+    assertEquals(5, groups.get(1).items.size());
+    assertEquals(6, groups.get(2).items.size());
   }
 
   public Cluster makeCluster() throws AmbariException {