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.cassandra.net;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import org.apache.cassandra.utils.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.Function;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.HashMultimap;
  import com.google.common.collect.SetMultimap;
  import org.apache.commons.lang3.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.cassandra.gms.EndpointState;
  import org.apache.cassandra.gms.Gossiper;
  import org.apache.cassandra.gms.IEndpointStateChangeSubscriber;
  import org.apache.cassandra.locator.InetAddressAndPort;
  import org.apache.cassandra.utils.FBUtilities;
  
  import static java.util.stream.Collectors.groupingBy;
  import static java.util.stream.Collectors.mapping;
  import static java.util.stream.Collectors.toList;
  import static org.apache.cassandra.net.Verb.PING_REQ;
  import static org.apache.cassandra.net.ConnectionType.LARGE_MESSAGES;
  import static org.apache.cassandra.net.ConnectionType.SMALL_MESSAGES;
  import static org.apache.cassandra.utils.Clock.Global.nanoTime;
  import static org.apache.cassandra.utils.concurrent.CountDownLatch.newCountDownLatch;
  
  public class StartupClusterConnectivityChecker
  {
      private static final Logger logger = LoggerFactory.getLogger(StartupClusterConnectivityChecker.class);
  
      private final boolean blockForRemoteDcs;
      private final long timeoutNanos;
  
      public static StartupClusterConnectivityChecker create(long timeoutSecs, boolean blockForRemoteDcs)
      {
          if (timeoutSecs > 100)
              logger.warn("setting the block-for-peers timeout (in seconds) to {} might be a bit excessive, but using it nonetheless", timeoutSecs);
          long timeoutNanos = TimeUnit.SECONDS.toNanos(timeoutSecs);
  
          return new StartupClusterConnectivityChecker(timeoutNanos, blockForRemoteDcs);
      }
  
      @VisibleForTesting
      StartupClusterConnectivityChecker(long timeoutNanos, boolean blockForRemoteDcs)
      {
          this.blockForRemoteDcs = blockForRemoteDcs;
          this.timeoutNanos = timeoutNanos;
      }
  
      /**
       * @param peers The currently known peers in the cluster; argument is not modified.
       * @param getDatacenterSource A function for mapping peers to their datacenter.
       * @return true if the requested percentage of peers are marked ALIVE in gossip and have their connections opened;
       * else false.
       */
      public boolean execute(Set<InetAddressAndPort> peers, Function<InetAddressAndPort, String> getDatacenterSource)
      {
          if (peers == null || this.timeoutNanos < 0)
              return true;
  
          // make a copy of the set, to avoid mucking with the input (in case it's a sensitive collection)
          peers = new HashSet<>(peers);
          InetAddressAndPort localAddress = FBUtilities.getBroadcastAddressAndPort();
          String localDc = getDatacenterSource.apply(localAddress);
  
          peers.remove(localAddress);
          if (peers.isEmpty())
              return true;
  
          // make a copy of the datacenter mapping (in case gossip updates happen during this method or some such)
          Map<InetAddressAndPort, String> peerToDatacenter = new HashMap<>();
          SetMultimap<String, InetAddressAndPort> datacenterToPeers = HashMultimap.create();
  
          for (InetAddressAndPort peer : peers)
          {
              String datacenter = getDatacenterSource.apply(peer);
              peerToDatacenter.put(peer, datacenter);
              datacenterToPeers.put(datacenter, peer);
          }
  
          // In the case where we do not want to block startup on remote datacenters (e.g. because clients only use
          // LOCAL_X consistency levels), we remove all other datacenter hosts from the mapping and we only wait
          // on the remaining local datacenter.
          if (!blockForRemoteDcs)
          {
              datacenterToPeers.keySet().retainAll(Collections.singleton(localDc));
              logger.info("Blocking coordination until only a single peer is DOWN in the local datacenter, timeout={}s",
                          TimeUnit.NANOSECONDS.toSeconds(timeoutNanos));
          }
          else
          {
              logger.info("Blocking coordination until only a single peer is DOWN in each datacenter, timeout={}s",
                          TimeUnit.NANOSECONDS.toSeconds(timeoutNanos));
          }
  
          // The threshold is 3 because for each peer we want to have 3 acks,
          // one for small message connection, one for large message connnection and one for alive event from gossip.
          AckMap acks = new AckMap(3, peers);
          Map<String, CountDownLatch> dcToRemainingPeers = new HashMap<>(datacenterToPeers.size());
          for (String datacenter: datacenterToPeers.keys())
          {
              dcToRemainingPeers.put(datacenter,
                                     newCountDownLatch(Math.max(datacenterToPeers.get(datacenter).size() - 1, 0)));
          }
  
          long startNanos = nanoTime();
  
          // set up a listener to react to new nodes becoming alive (in gossip), and account for all the nodes that are already alive
          Set<InetAddressAndPort> alivePeers = Collections.newSetFromMap(new ConcurrentHashMap<>());
          AliveListener listener = new AliveListener(alivePeers, dcToRemainingPeers, acks, peerToDatacenter::get);
          Gossiper.instance.register(listener);
  
          // send out a ping message to open up the non-gossip connections to all peers. Note that this sends the
          // ping messages to _all_ peers, not just the ones we block for in dcToRemainingPeers.
          sendPingMessages(peers, dcToRemainingPeers, acks, peerToDatacenter::get);
  
          for (InetAddressAndPort peer : peers)
          {
              if (Gossiper.instance.isAlive(peer) && alivePeers.add(peer) && acks.incrementAndCheck(peer))
              {
                  String datacenter = peerToDatacenter.get(peer);
                  // We have to check because we might only have the local DC in the map
                  if (dcToRemainingPeers.containsKey(datacenter))
                      dcToRemainingPeers.get(datacenter).decrement();
              }
          }
  
          boolean succeeded = true;
          for (CountDownLatch countDownLatch : dcToRemainingPeers.values())
          {
              long remainingNanos = Math.max(1, timeoutNanos - (nanoTime() - startNanos));
              //noinspection UnstableApiUsage
              succeeded &= countDownLatch.awaitUninterruptibly(remainingNanos, TimeUnit.NANOSECONDS);
          }
  
          Gossiper.instance.unregister(listener);
  
          if (succeeded)
          {
              logger.info("Ensured sufficient healthy connections with {} after {} milliseconds",
                          dcToRemainingPeers.keySet(), TimeUnit.NANOSECONDS.toMillis(nanoTime() - startNanos));
          }
          else
          {
              // dc -> missing peer host addresses
              Map<String, List<String>> peersDown = acks.getMissingPeers().stream()
                                                        .collect(groupingBy(peer -> {
                                                                                String dc = peerToDatacenter.get(peer);
                                                                                if (dc != null)
                                                                                    return dc;
                                                                                return StringUtils.defaultString(getDatacenterSource.apply(peer), "unknown");
                                                                            },
                                                                            mapping(InetAddressAndPort::getHostAddressAndPort,
                                                                                    toList())));
              logger.warn("Timed out after {} milliseconds, was waiting for remaining peers to connect: {}",
                          TimeUnit.NANOSECONDS.toMillis(nanoTime() - startNanos),
                          peersDown);
          }
  
          return succeeded;
      }
  
      /**
       * Sends a "connection warmup" message to each peer in the collection, on every {@link ConnectionType}
       * used for internode messaging (that is not gossip).
       */
      private void sendPingMessages(Set<InetAddressAndPort> peers, Map<String, CountDownLatch> dcToRemainingPeers,
                                    AckMap acks, Function<InetAddressAndPort, String> getDatacenter)
      {
          RequestCallback responseHandler = msg -> {
              if (acks.incrementAndCheck(msg.from()))
              {
                  String datacenter = getDatacenter.apply(msg.from());
                  // We have to check because we might only have the local DC in the map
                  if (dcToRemainingPeers.containsKey(datacenter))
                      dcToRemainingPeers.get(datacenter).decrement();
              }
          };
  
          Message<PingRequest> small = Message.out(PING_REQ, PingRequest.forSmall);
          Message<PingRequest> large = Message.out(PING_REQ, PingRequest.forLarge);
          for (InetAddressAndPort peer : peers)
          {
              MessagingService.instance().sendWithCallback(small, peer, responseHandler, SMALL_MESSAGES);
              MessagingService.instance().sendWithCallback(large, peer, responseHandler, LARGE_MESSAGES);
          }
      }
  
      /**
       * A trivial implementation of {@link IEndpointStateChangeSubscriber} that really only cares about
       * {@link #onAlive(InetAddressAndPort, EndpointState)} invocations.
       */
      private static final class AliveListener implements IEndpointStateChangeSubscriber
      {
          private final Map<String, CountDownLatch> dcToRemainingPeers;
          private final Set<InetAddressAndPort> livePeers;
          private final Function<InetAddressAndPort, String> getDatacenter;
          private final AckMap acks;
  
          AliveListener(Set<InetAddressAndPort> livePeers, Map<String, CountDownLatch> dcToRemainingPeers,
                        AckMap acks, Function<InetAddressAndPort, String> getDatacenter)
          {
              this.livePeers = livePeers;
              this.dcToRemainingPeers = dcToRemainingPeers;
              this.acks = acks;
              this.getDatacenter = getDatacenter;
          }
  
          public void onAlive(InetAddressAndPort endpoint, EndpointState state)
          {
              if (livePeers.add(endpoint) && acks.incrementAndCheck(endpoint))
              {
                  String datacenter = getDatacenter.apply(endpoint);
                  if (dcToRemainingPeers.containsKey(datacenter))
                      dcToRemainingPeers.get(datacenter).decrement();
              }
          }
      }
  
      private static final class AckMap
      {
          private final int threshold;
          private final Map<InetAddressAndPort, AtomicInteger> acks;
  
          AckMap(int threshold, Iterable<InetAddressAndPort> initialPeers)
          {
              this.threshold = threshold;
              acks = new ConcurrentHashMap<>();
              for (InetAddressAndPort peer : initialPeers)
                  initOrGetCounter(peer);
          }
  
          boolean incrementAndCheck(InetAddressAndPort address)
          {
              return initOrGetCounter(address).incrementAndGet() == threshold;
          }
  
          /**
           * Get a list of peers that has not fully ack'd, i.e. not reaching threshold acks
           */
          List<InetAddressAndPort> getMissingPeers()
          {
              List<InetAddressAndPort> missingPeers = new ArrayList<>();
              for (Map.Entry<InetAddressAndPort, AtomicInteger> entry : acks.entrySet())
              {
                  if (entry.getValue().get() < threshold)
                      missingPeers.add(entry.getKey());
              }
              return missingPeers;
          }
  
          // init the counter for the peer just in case
          private AtomicInteger initOrGetCounter(InetAddressAndPort address)
          {
              return acks.computeIfAbsent(address, addr -> new AtomicInteger(0));
          }
      }
  }
  