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.bookkeeper.stream.storage.impl.sc;
  
  import com.google.common.collect.Lists;
  import com.google.common.collect.Maps;
  import com.google.common.collect.Sets;
  import com.google.common.util.concurrent.UncheckedExecutionException;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.stream.Collectors;
  import java.util.stream.LongStream;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.stream.proto.cluster.ClusterAssignmentData;
  import org.apache.bookkeeper.stream.proto.cluster.ClusterMetadata;
  import org.apache.bookkeeper.stream.proto.cluster.ServerAssignmentData;
  import org.apache.commons.lang3.tuple.Pair;
  /**
   * The default implementation of storage container controller.
   *
   * <p>The goal of this controller is uniformly distributing storage containers across all alive servers in
   * the cluster.
   *
   * <p>The algorithm here is based on the count-based stream balancer in distributedlog-proxy-server.
   */
  @Slf4j
  public class DefaultStorageContainerController implements StorageContainerController {
  
      static final class ServerAssignmentDataComparator
          implements Comparator<Pair<BookieId, LinkedList<Long>>> {
  
          @Override
          public int compare(Pair<BookieId, LinkedList<Long>> o1,
                             Pair<BookieId, LinkedList<Long>> o2) {
              int res = Integer.compare(o1.getValue().size(), o2.getValue().size());
              if (0 == res) {
                  // two servers have same number of container
                  // the order of these two servers doesn't matter, so use any attribute than can provide deterministic
                  // ordering during state computation is good enough
                  return String.CASE_INSENSITIVE_ORDER.compare(
                      o1.getKey().toString(),
                      o2.getKey().toString());
              } else {
                  return res;
              }
          }
      }
  
      @Override
      public ClusterAssignmentData computeIdealState(ClusterMetadata clusterMetadata,
                                                     ClusterAssignmentData currentState,
                                                     Set<BookieId> currentCluster) {
  
          if (currentCluster.isEmpty()) {
              log.info("Current cluster is empty. No alive server is found.");
              return currentState;
          }
  
          // 1. get current server assignments
          Map<BookieId, Set<Long>> currentServerAssignments;
          try {
              currentServerAssignments = currentState.getServersMap()
                  .entrySet()
                  .stream()
                  .collect(Collectors.toMap(e1 -> {
                          return BookieId.parse(e1.getKey());
                      },
                      e2 -> e2.getValue().getContainersList().stream().collect(Collectors.toSet())
                  ));
          } catch (UncheckedExecutionException uee) {
              log.warn("Invalid cluster assignment data is found : {} - {}. Recompute assignment from empty state",
                  currentState, uee.getCause().getMessage());
              currentServerAssignments = Maps.newHashMap();
          }
          Set<BookieId> currentServersAssigned = currentServerAssignments.keySet();
  
          // 2. if no servers is assigned, initialize the ideal state
          if (currentServersAssigned.isEmpty()) {
              return initializeIdealState(clusterMetadata, currentCluster);
          }
  
          // 3. get the cluster diffs
          Set<BookieId> serversAdded =
              Sets.difference(currentCluster, currentServersAssigned).immutableCopy();
          Set<BookieId> serversRemoved =
              Sets.difference(currentServersAssigned, currentCluster).immutableCopy();
  
          if (serversAdded.isEmpty() && serversRemoved.isEmpty()) {
              // cluster is unchanged, assuming the current state is ideal, no re-assignment is required.
              return currentState;
          }
  
          log.info("Storage container controller detects cluster changed:\n"
                  + "\t {} servers added: {}\n\t {} servers removed: {}",
              serversAdded.size(), serversAdded, serversRemoved.size(), serversRemoved);
  
          // 4. compute the containers that owned by servers removed. these containers are needed to be reassigned.
          Set<Long> containersToReassign = currentServerAssignments.entrySet().stream()
              .filter(serverEntry -> !currentCluster.contains(serverEntry.getKey()))
              .flatMap(serverEntry -> serverEntry.getValue().stream())
              .collect(Collectors.toSet());
  
          // 5. use an ordered set as priority deque to sort the servers by the number of assigned containers
          TreeSet<Pair<BookieId, LinkedList<Long>>> assignmentQueue =
              new TreeSet<>(new ServerAssignmentDataComparator());
          for (Map.Entry<BookieId, Set<Long>> entry : currentServerAssignments.entrySet()) {
              BookieId host = entry.getKey();
  
              if (!currentCluster.contains(host)) {
                  if (log.isTraceEnabled()) {
                      log.trace("Host {} is not in current cluster anymore", host);
                  }
                  continue;
              } else {
                  if (log.isTraceEnabled()) {
                      log.trace("Adding host {} to assignment queue", host);
                  }
                  assignmentQueue.add(Pair.of(host, Lists.newLinkedList(entry.getValue())));
              }
          }
  
          // 6. add new servers
          for (BookieId server : serversAdded) {
              assignmentQueue.add(Pair.of(server, Lists.newLinkedList()));
          }
  
          // 7. assign the containers that are needed to be reassigned.
          for (Long containerId : containersToReassign) {
              Pair<BookieId, LinkedList<Long>> leastLoadedServer = assignmentQueue.pollFirst();
              leastLoadedServer.getValue().add(containerId);
              assignmentQueue.add(leastLoadedServer);
          }
  
          // 8. rebalance the containers if needed
          int diffAllowed;
          if (assignmentQueue.size() > clusterMetadata.getNumStorageContainers()) {
              diffAllowed = 1;
          } else {
              diffAllowed = clusterMetadata.getNumStorageContainers() % assignmentQueue.size() == 0 ? 0 : 1;
          }
  
          Pair<BookieId, LinkedList<Long>> leastLoaded = assignmentQueue.first();
          Pair<BookieId, LinkedList<Long>> mostLoaded = assignmentQueue.last();
          while (mostLoaded.getValue().size() - leastLoaded.getValue().size() > diffAllowed) {
              leastLoaded = assignmentQueue.pollFirst();
              mostLoaded = assignmentQueue.pollLast();
  
              // move container from mostLoaded to leastLoaded
              Long containerId = mostLoaded.getValue().removeFirst();
              // add the container to the end to avoid balancing this container again.
              leastLoaded.getValue().addLast(containerId);
  
              assignmentQueue.add(leastLoaded);
              assignmentQueue.add(mostLoaded);
  
              leastLoaded = assignmentQueue.first();
              mostLoaded = assignmentQueue.last();
          }
  
          // 9. the new ideal state is computed, finalize it
          Map<String, ServerAssignmentData> newAssignmentMap = Maps.newHashMap();
          assignmentQueue.forEach(assignment -> newAssignmentMap.put(
              assignment.getKey().toString(),
              ServerAssignmentData.newBuilder()
                  .addAllContainers(assignment.getValue())
                  .build()));
          return ClusterAssignmentData.newBuilder()
              .putAllServers(newAssignmentMap)
              .build();
      }
  
      static ClusterAssignmentData initializeIdealState(ClusterMetadata clusterMetadata,
                                                        Set<BookieId> currentCluster) {
          List<BookieId> serverList = Lists.newArrayListWithExpectedSize(currentCluster.size());
          serverList.addAll(currentCluster);
          Collections.shuffle(serverList);
  
          int numServers = currentCluster.size();
          int numTotalContainers = (int) clusterMetadata.getNumStorageContainers();
          int numContainersPerServer = numTotalContainers / currentCluster.size();
  
          Map<String, ServerAssignmentData> assignmentMap = Maps.newHashMap();
          for (int serverIdx = 0; serverIdx < serverList.size(); serverIdx++) {
              BookieId server = serverList.get(serverIdx);
  
              int finalServerIdx = serverIdx;
              ServerAssignmentData assignmentData = ServerAssignmentData.newBuilder()
                  .addAllContainers(
                      LongStream.rangeClosed(0, numContainersPerServer).boxed()
                          .map(j -> j * numServers + finalServerIdx)
                          .filter(containerId -> containerId < numTotalContainers)
                          .collect(Collectors.toSet()))
                  .build();
              assignmentMap.put(server.toString(), assignmentData);
          }
  
          return ClusterAssignmentData.newBuilder()
              .putAllServers(assignmentMap)
              .build();
      }
  
  }
  