/*-
 * -\-\-
 * Helios Services
 * --
 * Copyright (C) 2016 Spotify AB
 * --
 * Licensed 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 com.spotify.helios.master;

import static com.google.common.base.Optional.fromNullable;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.collect.Lists.newArrayList;
import static com.google.common.collect.Lists.reverse;
import static com.spotify.helios.common.descriptors.DeploymentGroup.RollingUpdateReason.HOSTS_CHANGED;
import static com.spotify.helios.common.descriptors.DeploymentGroup.RollingUpdateReason.MANUAL;
import static com.spotify.helios.common.descriptors.DeploymentGroupStatus.State.FAILED;
import static com.spotify.helios.common.descriptors.DeploymentGroupStatus.State.ROLLING_OUT;
import static com.spotify.helios.common.descriptors.Descriptor.parse;
import static com.spotify.helios.common.descriptors.HostStatus.Status.DOWN;
import static com.spotify.helios.common.descriptors.HostStatus.Status.UP;
import static com.spotify.helios.servicescommon.coordination.ZooKeeperOperations.check;
import static com.spotify.helios.servicescommon.coordination.ZooKeeperOperations.create;
import static com.spotify.helios.servicescommon.coordination.ZooKeeperOperations.delete;
import static com.spotify.helios.servicescommon.coordination.ZooKeeperOperations.set;
import static java.lang.Boolean.TRUE;
import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
import static java.util.Collections.singletonList;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonMappingException;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;
import com.spotify.helios.common.HeliosRuntimeException;
import com.spotify.helios.common.Json;
import com.spotify.helios.common.descriptors.AgentInfo;
import com.spotify.helios.common.descriptors.Deployment;
import com.spotify.helios.common.descriptors.DeploymentGroup;
import com.spotify.helios.common.descriptors.DeploymentGroupStatus;
import com.spotify.helios.common.descriptors.DeploymentGroupTasks;
import com.spotify.helios.common.descriptors.Goal;
import com.spotify.helios.common.descriptors.HostInfo;
import com.spotify.helios.common.descriptors.HostStatus;
import com.spotify.helios.common.descriptors.Job;
import com.spotify.helios.common.descriptors.JobId;
import com.spotify.helios.common.descriptors.JobStatus;
import com.spotify.helios.common.descriptors.PortMapping;
import com.spotify.helios.common.descriptors.RolloutOptions;
import com.spotify.helios.common.descriptors.RolloutTask;
import com.spotify.helios.common.descriptors.Task;
import com.spotify.helios.common.descriptors.TaskStatus;
import com.spotify.helios.common.descriptors.TaskStatusEvent;
import com.spotify.helios.common.descriptors.ThrottleState;
import com.spotify.helios.rollingupdate.DeploymentGroupEventFactory;
import com.spotify.helios.rollingupdate.RollingUndeployPlanner;
import com.spotify.helios.rollingupdate.RollingUpdateError;
import com.spotify.helios.rollingupdate.RollingUpdateOp;
import com.spotify.helios.rollingupdate.RollingUpdateOpFactory;
import com.spotify.helios.rollingupdate.RollingUpdatePlanner;
import com.spotify.helios.servicescommon.EventSender;
import com.spotify.helios.servicescommon.VersionedValue;
import com.spotify.helios.servicescommon.ZooKeeperRegistrarUtil;
import com.spotify.helios.servicescommon.coordination.Node;
import com.spotify.helios.servicescommon.coordination.Paths;
import com.spotify.helios.servicescommon.coordination.ZooKeeperClient;
import com.spotify.helios.servicescommon.coordination.ZooKeeperClientProvider;
import com.spotify.helios.servicescommon.coordination.ZooKeeperOperation;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.UUID;
import java.util.function.Predicate;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.KeeperException.BadVersionException;
import org.apache.zookeeper.KeeperException.NoNodeException;
import org.apache.zookeeper.KeeperException.NodeExistsException;
import org.apache.zookeeper.KeeperException.NotEmptyException;
import org.apache.zookeeper.OpResult;
import org.apache.zookeeper.data.Stat;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * The Helios Master's view into ZooKeeper.
 */
public class ZooKeeperMasterModel implements MasterModel {
  private static final Comparator<TaskStatusEvent> EVENT_COMPARATOR =
      new Comparator<TaskStatusEvent>() {
        @Override
        public int compare(TaskStatusEvent arg0, TaskStatusEvent arg1) {
          if (arg1.getTimestamp() > arg0.getTimestamp()) {
            return -1;
          } else if (arg1.getTimestamp() == arg0.getTimestamp()) {
            return 0;
          } else {
            return 1;
          }
        }
      };

  private static final Logger log = LoggerFactory.getLogger(ZooKeeperMasterModel.class);

  public static final Map<JobId, TaskStatus> EMPTY_STATUSES = emptyMap();
  public static final TypeReference<HostInfo>
      HOST_INFO_TYPE =
      new TypeReference<HostInfo>() {};
  public static final TypeReference<AgentInfo>
      AGENT_INFO_TYPE =
      new TypeReference<AgentInfo>() {};
  public static final TypeReference<Map<String, String>>
      STRING_MAP_TYPE =
      new TypeReference<Map<String, String>>() {};
  public static final TypeReference<List<String>>
      STRING_LIST_TYPE =
      new TypeReference<List<String>>() {};

  private static final DeploymentGroupEventFactory DEPLOYMENT_GROUP_EVENT_FACTORY =
      new DeploymentGroupEventFactory();

  private final ZooKeeperClientProvider provider;
  private final String name;
  private final List<EventSender> eventSenders;
  private final String deploymentGroupEventTopic;

  /**
   * Constructor.
   *
   * @param provider     {@link ZooKeeperClientProvider}
   * @param name         The hostname of the machine running the {@link MasterModel}
   * @param eventSenders {@link EventSender}
   */
  public ZooKeeperMasterModel(final ZooKeeperClientProvider provider,
                              final String name,
                              final List<EventSender> eventSenders,
                              final String deploymentGroupEventTopic) {
    this.provider = Preconditions.checkNotNull(provider);
    this.name = Preconditions.checkNotNull(name);
    this.eventSenders = Preconditions.checkNotNull(eventSenders);
    this.deploymentGroupEventTopic = deploymentGroupEventTopic;
  }

  /**
   * Registers a host into ZooKeeper.  The {@code id} is initially generated randomly by the Agent
   * and persisted on disk.  This way, in the event that you have two agents attempting to register
   * with the same value of @{code host}, the first one will win.
   */
  @Override
  public void registerHost(final String host, final String id) {
    final ZooKeeperClient client = provider.get("registerHost");
    try {
      ZooKeeperRegistrarUtil.registerHost(client, Paths.configHostId(host), host, id);
    } catch (Exception e) {
      throw new HeliosRuntimeException("registering host " + host + " failed", e);
    }
  }

  /**
   * Returns a list of the hosts/agents that have been registered.
   */
  @Override
  public List<String> listHosts() {
    try {
      // TODO (dano): only return hosts whose agents completed registration (i.e. has id nodes)
      return provider.get("listHosts").getChildren(Paths.configHosts());
    } catch (KeeperException.NoNodeException e) {
      return emptyList();
    } catch (KeeperException e) {
      throw new HeliosRuntimeException("listing hosts failed", e);
    }
  }

  @Override
  public List<String> listHosts(final String namePatternFilter) {
    Preconditions.checkNotNull(namePatternFilter, "namePatternFilter");
    final Predicate<String> matchesPattern = Pattern.compile(namePatternFilter).asPredicate();
    return listHosts().stream()
        .filter(matchesPattern)
        .collect(Collectors.toList());
  }

  /**
   * Returns a list of the host names of the currently running masters.
   */
  @Override
  public List<String> getRunningMasters() {
    final ZooKeeperClient client = provider.get("getRunningMasters");
    try {
      final List<String> masters = client.getChildren(Paths.statusMaster());
      final ImmutableList.Builder<String> upMasters = ImmutableList.builder();
      for (final String master : masters) {
        if (client.exists(Paths.statusMasterUp(master)) != null) {
          upMasters.add(master);
        }
      }
      return upMasters.build();
    } catch (KeeperException e) {
      throw new HeliosRuntimeException("listing masters failed", e);
    }
  }

  /**
   * Undoes the effect of {@link ZooKeeperMasterModel#registerHost(String, String)}.  Cleans up
   * any leftover host-related things.
   */
  @Override
  public void deregisterHost(final String host)
      throws HostNotFoundException, HostStillInUseException {
    final ZooKeeperClient client = provider.get("deregisterHost");
    ZooKeeperRegistrarUtil.deregisterHost(client, host);
  }

  /**
   * Adds a job into the configuration.
   */
  @Override
  public void addJob(final Job job) throws JobExistsException {
    log.info("adding job: {}", job);
    final JobId id = job.getId();
    final UUID operationId = UUID.randomUUID();
    final String creationPath = Paths.configJobCreation(id, operationId);
    final ZooKeeperClient client = provider.get("addJob");
    try {
      try {
        client.ensurePath(Paths.historyJob(id));
        client.transaction(create(Paths.configJob(id), job),
            create(Paths.configJobRefShort(id), id),
            create(Paths.configJobHosts(id)),
            create(creationPath),
            // Touch the jobs root node so that its version is bumped on every job
            // change down the tree. Effectively, make it that version == cVersion.
            set(Paths.configJobs(), UUID.randomUUID().toString().getBytes()));
      } catch (final NodeExistsException e) {
        if (client.exists(creationPath) != null) {
          // The job was created, we're done here
          return;
        }
        throw new JobExistsException(id.toString());
      }
    } catch (NoNodeException e) {
      throw new HeliosRuntimeException("adding job " + job + " failed due to missing ZK path: "
                                       + e.getPath(), e);
    } catch (final KeeperException e) {
      throw new HeliosRuntimeException("adding job " + job + " failed", e);
    }
  }

  /**
   * Given a jobId, returns the N most recent events in its history in the cluster.
   */
  @Override
  public List<TaskStatusEvent> getJobHistory(final JobId jobId) throws JobDoesNotExistException {
    return getJobHistory(jobId, null);
  }

  /**
   * Given a jobId and host, returns the N most recent events in its history on that host in the
   * cluster.
   */
  @Override
  public List<TaskStatusEvent> getJobHistory(final JobId jobId, final String host)
      throws JobDoesNotExistException {
    final Job descriptor = getJob(jobId);
    if (descriptor == null) {
      throw new JobDoesNotExistException(jobId);
    }
    final ZooKeeperClient client = provider.get("getJobHistory");
    final List<String> hosts;

    try {
      hosts = (!isNullOrEmpty(host)) ? singletonList(host) :
              client.getChildren(Paths.historyJobHosts(jobId));
    } catch (NoNodeException e) {
      return emptyList();
    } catch (KeeperException e) {
      throw new RuntimeException(e);
    }

    final List<TaskStatusEvent> jsEvents = Lists.newArrayList();

    for (final String h : hosts) {
      final List<String> events;
      try {
        events = client.getChildren(Paths.historyJobHostEvents(jobId, h));
      } catch (NoNodeException e) {
        continue;
      } catch (KeeperException e) {
        throw new RuntimeException(e);
      }

      for (final String event : events) {
        try {
          final byte[] data = client.getData(Paths.historyJobHostEventsTimestamp(
              jobId, h, Long.valueOf(event)));
          final TaskStatus status = Json.read(data, TaskStatus.class);
          jsEvents.add(new TaskStatusEvent(status, Long.valueOf(event), h));
        } catch (NoNodeException e) { // ignore, it went away before we read it
        } catch (KeeperException | IOException e) {
          throw new RuntimeException(e);
        }
      }
    }

    return Ordering.from(EVENT_COMPARATOR).sortedCopy(jsEvents);
  }

  /**
   * Create a deployment group.
   *
   * <p>If successful, the following ZK nodes will be created:
   * <ul>
   * <li>/config/deployment-groups/[group-name]</li>
   * <li>/status/deployment-groups/[group-name]</li>
   * <li>/status/deployment-groups/[group-name]/hosts</li>
   * </ul>
   * These nodes are guaranteed to exist until the DG is removed.
   *
   * <p>If the operation fails no ZK nodes will be created. If any of the nodes above already exist
   * the operation will fail.
   *
   * @throws DeploymentGroupExistsException If a DG with the same name already exists.
   */
  @Override
  public void addDeploymentGroup(final DeploymentGroup deploymentGroup)
      throws DeploymentGroupExistsException {
    log.info("adding deployment-group: {}", deploymentGroup);
    final ZooKeeperClient client = provider.get("addDeploymentGroup");

    try {
      try {
        client.ensurePath(Paths.configDeploymentGroups());
        client.ensurePath(Paths.statusDeploymentGroups());
        client.transaction(
            create(Paths.configDeploymentGroup(deploymentGroup.getName()), deploymentGroup),
            create(Paths.statusDeploymentGroup(deploymentGroup.getName())),
            create(Paths.statusDeploymentGroupHosts(deploymentGroup.getName()),
                Json.asBytesUnchecked(emptyList())),
            create(Paths.statusDeploymentGroupRemovedHosts(deploymentGroup.getName()),
                Json.asBytesUnchecked(emptyList()))
        );
      } catch (final NodeExistsException e) {
        throw new DeploymentGroupExistsException(deploymentGroup.getName());
      }
    } catch (final KeeperException e) {
      throw new HeliosRuntimeException("adding deployment-group " + deploymentGroup + " failed", e);
    }
  }

  @Override
  public DeploymentGroup getDeploymentGroup(final String name)
      throws DeploymentGroupDoesNotExistException {
    log.debug("getting deployment-group: {}", name);
    final ZooKeeperClient client = provider.get("getDeploymentGroup");
    return getDeploymentGroup(client, name);
  }

  private DeploymentGroup getDeploymentGroup(final ZooKeeperClient client, final String name)
      throws DeploymentGroupDoesNotExistException {
    try {
      final byte[] data = client.getData(Paths.configDeploymentGroup(name));
      return Json.read(data, DeploymentGroup.class);
    } catch (NoNodeException e) {
      throw new DeploymentGroupDoesNotExistException(name);
    } catch (KeeperException | IOException e) {
      throw new HeliosRuntimeException("getting deployment-group " + name + " failed", e);
    }
  }

  /**
   * Remove a deployment group.
   *
   * <p>If successful, all ZK nodes associated with the DG will be deleted. Specifically these
   * nodes are guaranteed to be non-existent after a successful remove (not all of them might exist
   * before, though):
   * <ul>
   * <li>/config/deployment-groups/[group-name]</li>
   * <li>/status/deployment-groups/[group-name]</li>
   * <li>/status/deployment-groups/[group-name]/hosts</li>
   * <li>/status/deployment-groups/[group-name]/removed</li>
   * <li>/status/deployment-group-tasks/[group-name]</li>
   * </ul>
   * If the operation fails no ZK nodes will be removed.
   *
   * @throws DeploymentGroupDoesNotExistException If the DG does not exist.
   */
  @Override
  public void removeDeploymentGroup(final String name) throws DeploymentGroupDoesNotExistException {
    log.info("removing deployment-group: name={}", name);
    final ZooKeeperClient client = provider.get("removeDeploymentGroup");
    try {
      client.ensurePath(Paths.configDeploymentGroups());
      client.ensurePath(Paths.statusDeploymentGroups());
      client.ensurePath(Paths.statusDeploymentGroupTasks());

      final List<ZooKeeperOperation> operations = Lists.newArrayList();

      final List<String> paths = ImmutableList.of(
          Paths.configDeploymentGroup(name),
          Paths.statusDeploymentGroup(name),
          Paths.statusDeploymentGroupHosts(name),
          Paths.statusDeploymentGroupRemovedHosts(name),
          Paths.statusDeploymentGroupTasks(name)
      );

      // For each deployment group path:
      //
      // * If it exists: delete it.
      // * If it doesn't exist, add and delete it in the same transaction. This is a round-about
      //   way of ensuring that it wasn't created when we commit the transaction.
      //
      // This is particularly important for /status/deployment-group-tasks/[group-name], which
      // might exist if a rolling-update is in progress. To avoid inconsistent state we make sure
      // it's deleted if it does exist.
      //
      // Having /status/deployment-group-tasks/[group-name] for removed groups around will cause
      // DGs to become slower and spam logs with errors so we want to avoid it.
      for (final String path : paths) {
        if (client.exists(path) == null) {
          operations.add(create(path));
        }
      }
      for (final String path : Lists.reverse(paths)) {
        operations.add(delete(path));
      }

      client.transaction(operations);
    } catch (final NoNodeException e) {
      throw new DeploymentGroupDoesNotExistException(name);
    } catch (final KeeperException e) {
      throw new HeliosRuntimeException("removing deployment-group " + name + " failed", e);
    }
  }

  /**
   * Determines whether we should allow deployment group hosts to be updated.
   *
   * <p>We don't want deployment groups that are ROLLING_OUT to change hosts in order to avoid the
   * following race:
   *
   * <p>- Manual rolling update A of a bad, broken job is triggered.
   * - The hosts change before rolling update A completes, triggering rolling update B.
   * - Rolling update B fails, because it inherited the bad, broken job from update A.
   * - The hosts change again, and we trigger unwanted rolling update C because the the previous
   *   rolling update (B) had reason=HOSTS_CHANGED.
   *
   * <p>It's safe to simply abort the hosts changed operation as updateOnHostsChange will be called
   * again by a reactor.
   *
   * @param status The status of the deployment group attempting to change hosts.
   *
   * @return True if it's safe to update the hosts.
   */
  private boolean allowHostChange(final DeploymentGroupStatus status) {
    if (status == null) {
      // We're in an unknown state. Go hog wild.
      return true;
    }
    return status.getState() != ROLLING_OUT;
  }

  /**
   * Determines whether we should trigger a rolling update when deployment group hosts change.
   *
   * <p>We want to avoid triggering an automatic rolling update if the most recent rolling update
   * was triggered manually, and failed.
   *
   * @param group  The deployment group that is changing hosts.
   * @param status The status of the aforementioned deployment group.
   *
   * @return True if we should perform a rolling update.
   */
  private boolean updateOnHostChange(final DeploymentGroup group,
                                     final DeploymentGroupStatus status) {
    if (status == null) {
      // We're in an unknown state. Go hog wild.
      return true;
    }

    if (group.getRollingUpdateReason() == null) {
      // The last rolling update didn't fail, so there's no reason to expect this one will.
      return status.getState() != FAILED;
    }

    if (group.getRollingUpdateReason() == HOSTS_CHANGED && status.getState() == FAILED) {
      // The last rolling update failed, but it was triggered by hosts changing. Try again.
      return true;
    }

    // The last rolling update didn't fail, so there's no reason to expect this one will.
    return status.getState() != FAILED;
  }

  private List<String> removedHosts(final List<String> currentHosts,
                                    final List<String> newHosts,
                                    final List<String> previouslyRemovedHosts) {
    final Set<String> ch = ImmutableSet.copyOf(currentHosts);
    final Set<String> nh = ImmutableSet.copyOf(newHosts);
    final Set<String> prh = ImmutableSet.copyOf(previouslyRemovedHosts);

    // Calculate the freshly removed hosts (current - new) and add in any previously removed hosts
    // that haven't been undeployed yet.
    final Set<String> rh = Sets.union(Sets.difference(ch, nh), prh);

    // Finally, only include hosts that are still registered
    return ImmutableList.copyOf(Sets.intersection(rh, ImmutableSet.copyOf(listHosts())));
  }

  @Override
  public void updateDeploymentGroupHosts(final String groupName, final List<String> hosts)
      throws DeploymentGroupDoesNotExistException {
    log.debug("updating deployment-group hosts: name={}", groupName);
    final ZooKeeperClient client = provider.get("updateDeploymentGroupHosts");
    try {
      final DeploymentGroupStatus status = getDeploymentGroupStatus(groupName);
      if (!allowHostChange(status)) {
        return;
      }

      // statusDeploymentGroupRemovedHosts may not exist for deployment groups created before it was
      // introduced.
      client.ensurePathAndSetData(Paths.statusDeploymentGroupRemovedHosts(groupName),
          Json.asBytesUnchecked(emptyList()));

      final List<String> curHosts = getHosts(client, Paths.statusDeploymentGroupHosts(groupName));
      final List<String> previouslyRemovedHosts = getHosts(
          client, Paths.statusDeploymentGroupRemovedHosts(groupName));

      final List<String> removedHosts = removedHosts(curHosts, hosts, previouslyRemovedHosts);

      if (hosts.equals(curHosts) && removedHosts.equals(previouslyRemovedHosts)) {
        return;
      }

      log.info("for deployment-group name={}, curHosts={}, new hosts={}, "
               + "previouslyRemovedHosts={}, derived removedHosts={}",
          groupName, curHosts, hosts, previouslyRemovedHosts, removedHosts);

      final List<ZooKeeperOperation> ops = Lists.newArrayList();
      ops.add(set(Paths.statusDeploymentGroupHosts(groupName), Json.asBytes(hosts)));
      ops.add(set(Paths.statusDeploymentGroupRemovedHosts(groupName), Json.asBytes(removedHosts)));

      final Node dgn = client.getNode(Paths.configDeploymentGroup(groupName));
      final Integer deploymentGroupVersion = dgn.getStat().getVersion();
      DeploymentGroup deploymentGroup = Json.read(dgn.getBytes(), DeploymentGroup.class);

      List<Map<String, Object>> events = ImmutableList.of();

      if (deploymentGroup.getJobId() != null && updateOnHostChange(deploymentGroup, status)) {
        deploymentGroup = deploymentGroup.toBuilder()
            .setRollingUpdateReason(HOSTS_CHANGED)
            .build();

        // Fail transaction if the deployment group has been updated elsewhere.
        ops.add(check(Paths.configDeploymentGroup(groupName), deploymentGroupVersion));

        // NOTE: If the DG was removed this set() cause the transaction to fail, because
        // removing the DG removes this node. It's *important* that there's an operation that
        // causes the transaction to fail if the DG was removed or we'll end up with
        // inconsistent state.
        ops.add(set(Paths.configDeploymentGroup(deploymentGroup.getName()), deploymentGroup));

        final RollingUpdateOp op = getInitRollingUpdateOps(
            deploymentGroup, hosts, removedHosts, client);
        ops.addAll(op.operations());
        events = op.events();
      }

      log.info("starting zookeeper transaction for updateDeploymentGroupHosts on deployment-group: "
               + "name={} jobId={} operations={}",
          groupName, deploymentGroup.getJobId(), ops);

      client.transaction(ops);
      emitEvents(deploymentGroupEventTopic, events);
    } catch (BadVersionException e) {
      // some other master beat us in processing this host update. not exceptional.
      // ideally we would check the path in the exception, but curator doesn't provide a path
      // for exceptions thrown as part of a transaction.
      log.info("zookeeper transaction for updateDeploymentGroupHosts on deployment-group was "
               + "processed by another master: name={}", groupName);
    } catch (NoNodeException e) {
      throw new DeploymentGroupDoesNotExistException(groupName, e);
    } catch (KeeperException | IOException e) {
      throw new HeliosRuntimeException("updating deployment group hosts failed", e);
    }
  }

  /**
   * Returns a {@link RolloutOptions} instance that will replace null attributes in options with
   * values from two tiers of fallbacks. First fallback to job then
   * {@link RolloutOptions#getDefault()}.
   */
  static RolloutOptions rolloutOptionsWithFallback(final RolloutOptions options, final Job job) {
    return options.withFallback(
        job.getRolloutOptions() == null
        ? RolloutOptions.getDefault()
        : job.getRolloutOptions().withFallback(RolloutOptions.getDefault()));
  }

  @Override
  public void rollingUpdate(final DeploymentGroup deploymentGroup,
                            final JobId jobId,
                            final RolloutOptions options)
      throws DeploymentGroupDoesNotExistException, JobDoesNotExistException {
    checkNotNull(deploymentGroup, "deploymentGroup");

    final Job job = getJob(jobId);
    if (job == null) {
      throw new JobDoesNotExistException(jobId);
    }

    final RolloutOptions rolloutOptionsWithFallback = rolloutOptionsWithFallback(options, job);

    log.info("preparing to initiate rolling-update on deployment-group: "
             + "name={}, jobId={}, options={}",
        deploymentGroup.getName(), jobId, rolloutOptionsWithFallback);


    final DeploymentGroup updated = deploymentGroup.toBuilder()
        .setJobId(jobId)
        .setRolloutOptions(rolloutOptionsWithFallback)
        .setRollingUpdateReason(MANUAL)
        .build();

    final List<ZooKeeperOperation> operations = Lists.newArrayList();
    final ZooKeeperClient client = provider.get("rollingUpdate");

    operations.add(set(Paths.configDeploymentGroup(updated.getName()), updated));

    try {
      final RollingUpdateOp op = getInitRollingUpdateOps(updated, client);
      operations.addAll(op.operations());

      log.info("starting zookeeper transaction for rolling-update on "
               + "deployment-group name={} jobId={}. List of operations: {}",
          deploymentGroup.getName(), jobId, operations);

      client.transaction(operations);

      emitEvents(deploymentGroupEventTopic, op.events());
      log.info("initiated rolling-update on deployment-group: name={}, jobId={}",
          deploymentGroup.getName(), jobId);
    } catch (final NoNodeException e) {
      throw new DeploymentGroupDoesNotExistException(deploymentGroup.getName());
    } catch (final KeeperException e) {
      throw new HeliosRuntimeException(
          "rolling-update on deployment-group " + deploymentGroup.getName() + " failed. "
          + e.getMessage(), e);
    }
  }

  private RollingUpdateOp getInitRollingUpdateOps(final DeploymentGroup deploymentGroup,
                                                  final ZooKeeperClient zooKeeperClient)
      throws DeploymentGroupDoesNotExistException, KeeperException {
    final List<String> hosts = getDeploymentGroupHosts(deploymentGroup.getName());
    return getInitRollingUpdateOps(deploymentGroup, hosts, ImmutableList.of(), zooKeeperClient);
  }

  private RollingUpdateOp getInitRollingUpdateOps(final DeploymentGroup deploymentGroup,
                                                  final List<String> updateHosts,
                                                  final List<String> undeployHosts,
                                                  final ZooKeeperClient zooKeeperClient)
      throws KeeperException {
    final List<RolloutTask> rolloutTasks = new ArrayList<>();

    // give precedence to the updateHosts list so we don't end up in a state where we updated a host
    // and then removed the job from it (because of buggy logic in the calling method)
    final List<String> updateHostsCopy = new ArrayList<>(updateHosts);
    final List<String> undeployHostsCopy = new ArrayList<>(undeployHosts);
    undeployHostsCopy.removeAll(updateHostsCopy);

    // we only care about hosts that are UP
    final List<String> upHostsToUndeploy = undeployHostsCopy.stream()
        .filter(host -> checkHostUp(zooKeeperClient, host))
        .collect(Collectors.toList());
    final List<String> upHostsToDeploy = updateHostsCopy.stream()
        .filter(host -> checkHostUp(zooKeeperClient, host))
        .collect(Collectors.toList());

    rolloutTasks.addAll(RollingUndeployPlanner.of(deploymentGroup)
        .plan(upHostsToUndeploy));
    rolloutTasks.addAll(RollingUpdatePlanner.of(deploymentGroup)
        .plan(upHostsToDeploy));

    log.info("generated rolloutTasks for deployment-group name={} "
             + "updateHosts={} undeployHosts={}: {}",
        deploymentGroup.getName(), updateHosts, undeployHosts, rolloutTasks);

    final DeploymentGroupTasks tasks = DeploymentGroupTasks.newBuilder()
        .setRolloutTasks(rolloutTasks)
        .setTaskIndex(0)
        .setDeploymentGroup(deploymentGroup)
        .build();

    return new RollingUpdateOpFactory(tasks, DEPLOYMENT_GROUP_EVENT_FACTORY)
        .start(deploymentGroup, zooKeeperClient);
  }

  private Map<String, VersionedValue<DeploymentGroupTasks>> getDeploymentGroupTasks(
      final ZooKeeperClient client) {
    final String folder = Paths.statusDeploymentGroupTasks();
    try {
      final List<String> names;
      try {
        names = client.getChildren(folder);
      } catch (NoNodeException e) {
        return Collections.emptyMap();
      }

      final Map<String, VersionedValue<DeploymentGroupTasks>> ret = Maps.newHashMap();
      for (final String name : names) {
        final String path = Paths.statusDeploymentGroupTasks(name);
        try {
          final Node node = client.getNode(path);
          final byte[] data = node.getBytes();
          final int version = node.getStat().getVersion();
          if (data.length == 0) {
            // This can happen because of ensurePath creates an empty node
            log.debug("Ignoring empty deployment group tasks {}", name);
          } else {
            final DeploymentGroupTasks val = parse(data, DeploymentGroupTasks.class);
            ret.put(name, VersionedValue.of(val, version));
          }
        } catch (NoNodeException e) {
          // Ignore, the deployment group was deleted before we had a chance to read it.
          log.debug("Ignoring deleted deployment group tasks {}", name);
        }
      }
      return ret;
    } catch (KeeperException | IOException e) {
      throw new HeliosRuntimeException("getting deployment group tasks failed", e);
    }
  }

  private RollingUpdateOp processRollingUpdateTask(final ZooKeeperClient client,
