content: /*
   * 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.stack.upgrade;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.LinkedHashSet;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.stream.Stream;
  
  import javax.xml.bind.annotation.XmlElement;
  import javax.xml.bind.annotation.XmlType;
  
  import org.apache.ambari.server.stack.HostsType;
  import org.apache.ambari.server.stack.upgrade.UpgradePack.ProcessingComponent;
  import org.apache.ambari.server.stack.upgrade.orchestrate.StageWrapper;
  import org.apache.ambari.server.stack.upgrade.orchestrate.StageWrapper.Type;
  import org.apache.ambari.server.stack.upgrade.orchestrate.StageWrapperBuilder;
  import org.apache.ambari.server.stack.upgrade.orchestrate.TaskWrapper;
  import org.apache.ambari.server.stack.upgrade.orchestrate.TaskWrapperBuilder;
  import org.apache.ambari.server.stack.upgrade.orchestrate.UpgradeContext;
  import org.apache.ambari.server.utils.SetUtils;
  import org.apache.ambari.spi.upgrade.OrchestrationOptions;
  import org.apache.commons.collections4.CollectionUtils;
  import org.apache.commons.lang.ArrayUtils;
  import org.apache.commons.lang.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.google.common.base.Predicate;
  import com.google.common.collect.Collections2;
  import com.google.gson.JsonArray;
  import com.google.gson.JsonObject;
  import com.google.gson.JsonPrimitive;
  
  /**
   * Used for co-located services grouped together.
   */
  @XmlType(name="colocated")
  public class ColocatedGrouping extends Grouping {
  
    private static final Logger LOG = LoggerFactory.getLogger(ColocatedGrouping.class);
  
    @XmlElement(name="batch")
    public Batch batch;
  
  
    /**
     * {@inheritDoc}
     */
    @Override
    public StageWrapperBuilder getBuilder() {
      return new MultiHomedBuilder(this, batch, performServiceCheck, parallelScheduler);
    }
  
    private static class MultiHomedBuilder extends StageWrapperBuilder {
  
      private Batch m_batch;
      private boolean m_serviceCheck = true;
  
      // lists of tasks
      private List<TaskProxy> initialBatch = new LinkedList<>();
      private List<TaskProxy> finalBatches = new LinkedList<>();
  
  
      private MultiHomedBuilder(Grouping grouping, Batch batch,
          boolean serviceCheck, ParallelScheduler parallel) {
        super(grouping);
  
        m_batch = batch;
        m_serviceCheck = serviceCheck;
      }
  
      @Override
      public void add(UpgradeContext context, HostsType hostsType, String service,
          boolean clientOnly, ProcessingComponent pc, Map<String, String> params) {
  
        // !!! the percent is the number of co-located items that should be run
        // before pausing. This value has no bearing on how many should be in a single stage.
        int count = Double.valueOf(Math.ceil(
            (double) m_batch.percent / 100 * hostsType.getHosts().size())).intValue();
  
        LinkedHashSet<String> first = new LinkedHashSet<>();
        LinkedHashSet<String> remaining = new LinkedHashSet<>();
        hostsType.getHosts().stream().forEach(hostName -> {
          if (first.size() < count) {
            first.add(hostName);
          } else {
            remaining.add(hostName);
          }
        });
  
        // !!! resolve tasks only once
        List<Task> preTasks = resolveTasks(context, true, pc);
        Task task = resolveTask(context, pc);
        List<Task> postTasks = resolveTasks(context, false, pc);
        AtomicBoolean processInitial = new AtomicBoolean(true);
  
        int parallelCount = -1;
        OrchestrationOptions options = context.getOrchestrationOptions();
        if (null != options) {
          parallelCount = context.getOrchestrationOptions().getConcurrencyCount(
              context.getCluster().buildClusterInformation(), service, pc.name);
        }
  
        if (parallelCount < 1) {
          parallelCount = getParallelHostCount(context, 1);
        }
  
        // !!! stupid effective-final check
        final int hostCount = parallelCount;
  
        Stream.of(first, remaining).forEach(hosts -> {
  
          List<TaskProxy> targetList = processInitial.get() ? initialBatch : finalBatches;
  
          List<Set<String>> hostSplit = SetUtils.split(hosts, hostCount);
  
          hostSplit.forEach(hostSet -> {
  
            List<Task> tasks = preTasks;
  
            TaskProxy proxy;
            if (CollectionUtils.isNotEmpty(preTasks)) {
              // Our assumption is that all of the tasks in the StageWrapper are of
              // the same type.
              StageWrapper.Type type = preTasks.get(0).getStageWrapperType();
  
              proxy = new TaskProxy();
              proxy.clientOnly = clientOnly;
              proxy.message = getStageText("Preparing",
                  context.getComponentDisplay(service, pc.name), hostSet);
              proxy.tasks.addAll(TaskWrapperBuilder.getTaskList(service, pc.name,
                  HostsType.normal(new LinkedHashSet<>(hostSet)), tasks, params));
              proxy.service = service;
              proxy.component = pc.name;
              proxy.type = type;
  
              targetList.add(proxy);
            }
  
            if (null != task && RestartTask.class.isInstance(task)) {
              proxy = new TaskProxy();
              proxy.clientOnly = clientOnly;
              proxy.tasks.add(new TaskWrapper(service, pc.name, hostSet, params, task));
              proxy.restart = true;
              proxy.service = service;
              proxy.component = pc.name;
              proxy.type = Type.RESTART;
              proxy.message = getStageText("Restarting",
                  context.getComponentDisplay(service, pc.name), hostSet);
  
              targetList.add(proxy);
            }
  
            tasks = postTasks;
            if (CollectionUtils.isNotEmpty(preTasks)) {
              // Our assumption is that all of the tasks in the StageWrapper are of
              // the same type.
              StageWrapper.Type type = preTasks.get(0).getStageWrapperType();
  
              proxy = new TaskProxy();
              proxy.clientOnly = clientOnly;
              proxy.message = getStageText("Completing",
                  context.getComponentDisplay(service, pc.name), hostSet);
              proxy.tasks.addAll(TaskWrapperBuilder.getTaskList(service, pc.name,
                  HostsType.normal(new LinkedHashSet<>(hostSet)), tasks, params));
              proxy.service = service;
              proxy.component = pc.name;
              proxy.type = type;
  
              targetList.add(proxy);
            }
          });
  
          processInitial.set(false);
        });
  
      }
  
  
      /**
       * {@inheritDoc}
       */
      @Override
      public List<StageWrapper> build(UpgradeContext upgradeContext, List<StageWrapper> stageWrappers) {
  
        final List<Task> visitedServerSideTasks = new ArrayList<>();
  
        // !!! predicate to ensure server-side tasks are executed once only per grouping
        Predicate<Task> predicate = new Predicate<Task>() {
          @Override
          public boolean apply(Task input) {
            if (visitedServerSideTasks.contains(input)) {
              return false;
            }
  
            if (input.getType().isServerAction()) {
              visitedServerSideTasks.add(input);
            }
  
            return true;
          };
        };
  
        List<StageWrapper> results = new ArrayList<>(stageWrappers);
  
        if (LOG.isDebugEnabled()) {
          LOG.debug("RU initial: {}", initialBatch);
          LOG.debug("RU final: {}", finalBatches);
        }
  
        List<StageWrapper> befores = fromProxies(upgradeContext.getDirection(), initialBatch, predicate);
        results.addAll(befores);
  
        if (!befores.isEmpty()) {
  
          ManualTask task = new ManualTask();
          task.summary = m_batch.summary;
          List<String> messages = new ArrayList<>();
          messages.add(m_batch.message);
          task.messages = messages;
          formatFirstBatch(upgradeContext, task, befores);
  
          StageWrapper wrapper = new StageWrapper(
              StageWrapper.Type.SERVER_SIDE_ACTION,
              "Validate Partial " + upgradeContext.getDirection().getText(true),
              new TaskWrapper(null, null, Collections.emptySet(), task));
  
          results.add(wrapper);
        }
  
        results.addAll(fromProxies(upgradeContext.getDirection(), finalBatches, predicate));
  
        return results;
      }
  
      private List<StageWrapper> fromProxies(Direction direction,
          List<TaskProxy> proxies, Predicate<Task> predicate) {
  
        List<StageWrapper> results = new ArrayList<>();
  
        Set<String> serviceChecks = new HashSet<>();
  
  
        proxies.forEach(proxy -> {
          StageWrapper wrapper = null;
          List<StageWrapper> execwrappers = new ArrayList<>();
  
          if (!proxy.clientOnly) {
            serviceChecks.add(proxy.service);
          }
  
          if (!proxy.restart) {
            TaskWrapper[] tasks = proxy.getTasksArray(predicate);
  
            if (LOG.isDebugEnabled()) {
              for (TaskWrapper tw : tasks) {
                LOG.debug("{}", tw);
              }
            }
  
            if (ArrayUtils.isNotEmpty(tasks)) {
              wrapper = new StageWrapper(proxy.type, proxy.message, tasks);
            }
          } else {
            TaskWrapper[] tasks = proxy.getTasksArray(null);
  
            if (LOG.isDebugEnabled()) {
              for (TaskWrapper tw : tasks) {
                LOG.debug("{}", tw);
              }
            }
            // !!! TODO check parallelism values
            execwrappers.add(new StageWrapper(StageWrapper.Type.RESTART, proxy.message, tasks));
          }
  
          if (null != wrapper) {
            results.add(wrapper);
          }
  
          if (execwrappers.size() > 0) {
            results.addAll(execwrappers);
          }
        });
  
        if (direction.isUpgrade() && m_serviceCheck && serviceChecks.size() > 0) {
          // !!! add the service check task
          List<TaskWrapper> tasks = new ArrayList<>();
          Set<String> displays = new HashSet<>();
          for (String service : serviceChecks) {
            tasks.add(new TaskWrapper(service, "", Collections.emptySet(), new ServiceCheckTask()));
            displays.add(service);
          }
  
          StageWrapper wrapper = new StageWrapper(
              StageWrapper.Type.SERVICE_CHECK,
              "Service Check " + StringUtils.join(displays, ", "),
              tasks.toArray(new TaskWrapper[tasks.size()]));
  
          results.add(wrapper);
        }
  
        return results;
      }
  
      /**
       * Formats the first batch's text and adds json for use if needed.
       * @param ctx       the upgrade context to load component display names
       * @param task      the manual task representing the verification message
       * @param wrappers  the list of stage wrappers
       */
      private void formatFirstBatch(UpgradeContext ctx, ManualTask task, List<StageWrapper> wrappers) {
        Set<String> names = new LinkedHashSet<>();
        Map<String, Set<String>> compLocations = new HashMap<>();
  
        for (StageWrapper sw : wrappers) {
          for (TaskWrapper tw : sw.getTasks()) {
            if (StringUtils.isNotEmpty(tw.getService()) &&
                StringUtils.isNotBlank(tw.getComponent())) {
  
              for (String host : tw.getHosts()) {
                if (!compLocations.containsKey(host)) {
                  compLocations.put(host, new HashSet<>());
                }
                compLocations.get(host).add(tw.getComponent());
              }
  
              names.add(ctx.getComponentDisplay(
                  tw.getService(), tw.getComponent()));
            }
          }
        }
  
        for(int i = 0; i < task.messages.size(); i++){
          String message = task.messages.get(i);
          // !!! add the display names to the message, if needed
          if (message.contains("{{components}}")) {
            StringBuilder sb = new StringBuilder();
  
            List<String> compNames = new ArrayList<>(names);
  
            if (compNames.size() == 1) {
              sb.append(compNames.get(0));
            } else if (names.size() > 1) {
              String last = compNames.remove(compNames.size() - 1);
              sb.append(StringUtils.join(compNames, ", "));
              sb.append(" and ").append(last);
            }
  
            message = message.replace("{{components}}", sb.toString());
  
            //Add the updated message back to the message list.
            task.messages.set(i, message);
          }
        }
  
        // !!! build the structured out to attach to the manual task
        JsonArray arr = new JsonArray();
        for (Entry<String, Set<String>> entry : compLocations.entrySet()) {
          JsonObject obj = new JsonObject();
          obj.addProperty("host_name", entry.getKey());
  
          JsonArray comps = new JsonArray();
          for (String comp : entry.getValue()) {
            comps.add(new JsonPrimitive(comp));
          }
          obj.add("components", comps);
  
          arr.add(obj);
        }
  
        JsonObject master = new JsonObject();
        master.add("topology", arr);
  
        task.structuredOut = master.toString();
      }
    }
  
    /**
     * Represents all the tasks that need to be run for a host
     */
    private static class TaskProxy {
      private boolean restart = false;
      private String service;
      private String component;
      private String message;
      private Type type;
      private boolean clientOnly = false;
      private List<TaskWrapper> tasks = new ArrayList<>();
  
      @Override
      public String toString() {
        String s = "";
        for (TaskWrapper t : tasks) {
          s += component + "/" + t.getTasks() + " ";
        }
  
        return s;
      }
  
      /**
       * Get the task wrappers for this proxy.  Server-side tasks cannot be executed more than
       * one time per grouping.
       * @param predicate the predicate to determine if a server-side task has already been added to a wrapper.
       * @return the wrappers for a stage
       */
      private TaskWrapper[] getTasksArray(Predicate<Task> predicate) {
        if (null == predicate) {
          return tasks.toArray(new TaskWrapper[tasks.size()]);
        }
  
        List<TaskWrapper> interim = new ArrayList<>();
  
        for (TaskWrapper wrapper : tasks) {
          Collection<Task> filtered = Collections2.filter(wrapper.getTasks(), predicate);
  
          if (CollectionUtils.isNotEmpty(filtered)) {
            interim.add(wrapper);
          }
        }
  
        return interim.toArray(new TaskWrapper[interim.size()]);
      }
    }
  
  }
  