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.client;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import org.apache.bookkeeper.client.WeightedRandomSelection.WeightedObject;
  import org.apache.bookkeeper.conf.ClientConfiguration;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.proto.BookieClient;
  import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.GetBookieInfoCallback;
  import org.apache.bookkeeper.proto.BookkeeperProtocol;
  import org.apache.commons.collections4.CollectionUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A utility class to read {@link BookieInfo} from bookies.
   *
   * <p>NOTE: This class is tended to be used by this project only. External users should not rely on it directly.
   */
  public class BookieInfoReader {
      private static final Logger LOG = LoggerFactory.getLogger(BookieInfoReader.class);
      private static final long GET_BOOKIE_INFO_REQUEST_FLAGS =
          BookkeeperProtocol.GetBookieInfoRequest.Flags.TOTAL_DISK_CAPACITY_VALUE
                                 | BookkeeperProtocol.GetBookieInfoRequest.Flags.FREE_DISK_SPACE_VALUE;
  
      private final ScheduledExecutorService scheduler;
      private final BookKeeper bk;
      private final ClientConfiguration conf;
  
      /**
       * A class represents the information (e.g. disk usage, load) of a bookie.
       *
       * <p>NOTE: This class is tended to be used by this project only. External users should not rely on it directly.
       */
      public static class BookieInfo implements WeightedObject {
          private final long freeDiskSpace;
          private final long totalDiskSpace;
          public BookieInfo() {
              this(0L, 0L);
          }
          public BookieInfo(long totalDiskSpace, long freeDiskSpace) {
              this.totalDiskSpace = totalDiskSpace;
              this.freeDiskSpace = freeDiskSpace;
          }
          public long getFreeDiskSpace() {
              return freeDiskSpace;
          }
          public long getTotalDiskSpace() {
              return totalDiskSpace;
          }
          @Override
          public long getWeight() {
              return freeDiskSpace;
          }
          @Override
          public String toString() {
              return "FreeDiskSpace: " + this.freeDiskSpace + " TotalDiskCapacity: " + this.totalDiskSpace;
          }
      }
  
  
      /**
       * Tracks the most recently reported set of bookies from BookieWatcher as well
       * as current BookieInfo for bookies we've successfully queried.
       */
      private static class BookieInfoMap {
          /**
           * Contains the most recently obtained information on the contained bookies.
           * When an error happens querying a bookie, the entry is removed.
           */
          private final Map<BookieId, BookieInfo> infoMap = new HashMap<>();
  
          /**
           * Contains the most recently reported set of bookies from BookieWatcher
           * A partial query consists of every member of mostRecentlyReportedBookies
           * minus the entries in bookieInfoMap.
           */
          private Collection<BookieId> mostRecentlyReportedBookies = new ArrayList<>();
  
          public void updateBookies(Collection<BookieId> updatedBookieSet) {
              if (LOG.isDebugEnabled()) {
                  LOG.debug(
                          "updateBookies: current: {}, new: {}",
                          mostRecentlyReportedBookies, updatedBookieSet);
              }
              infoMap.keySet().retainAll(updatedBookieSet);
              mostRecentlyReportedBookies = updatedBookieSet;
          }
  
          @SuppressWarnings("unchecked")
          public Collection<BookieId> getPartialScanTargets() {
              return CollectionUtils.subtract(mostRecentlyReportedBookies, infoMap.keySet());
          }
  
          public Collection<BookieId> getFullScanTargets() {
              return mostRecentlyReportedBookies;
          }
  
          /**
           * Returns info for bookie, null if not known.
           *
           * @param bookie bookie for which to get info
           * @return Info for bookie, null otherwise
           */
          public BookieInfo getInfo(BookieId bookie) {
              return infoMap.get(bookie);
          }
  
          /**
           * Removes bookie from bookieInfoMap.
           *
           * @param bookie bookie on which we observed an error
           */
          public void clearInfo(BookieId bookie) {
              infoMap.remove(bookie);
          }
  
          /**
           * Report new info on bookie.
           *
           * @param bookie bookie for which we obtained new info
           * @param info the new info
           */
          public void gotInfo(BookieId bookie, BookieInfo info) {
              infoMap.put(bookie, info);
          }
  
          /**
           * Get bookie info map.
           */
          public Map<BookieId, BookieInfo> getBookieMap() {
              return infoMap;
          }
      }
      private final BookieInfoMap bookieInfoMap = new BookieInfoMap();
  
      /**
       * Tracks whether there is an execution in progress as well as whether
       * another is pending.
       */
      public enum State { UNQUEUED, PARTIAL, FULL }
      private static class InstanceState {
          private boolean running = false;
          private State queuedType = State.UNQUEUED;
  
          private boolean shouldStart() {
              if (running) {
                  return false;
              } else {
                  running = true;
                  return true;
              }
          }
  
          /**
           * Mark pending operation FULL and return true if there is no in-progress operation.
           *
           * @return True if we should execute a scan, False if there is already one running
           */
          public boolean tryStartFull() {
              queuedType = State.FULL;
              return shouldStart();
          }
  
          /**
           * Mark pending operation PARTIAL if not full and return true if there is no in-progress operation.
           *
           * @return True if we should execute a scan, False if there is already one running
           */
          public boolean tryStartPartial() {
              if (queuedType == State.UNQUEUED) {
                  queuedType = State.PARTIAL;
              }
              return shouldStart();
          }
  
          /**
           * Gets and clears queuedType.
           */
          public State getAndClearQueuedType() {
              State ret = queuedType;
              queuedType = State.UNQUEUED;
              return ret;
          }
  
          /**
           * If queuedType != UNQUEUED, returns true, leaves running equal to true
           * Otherwise, returns false and sets running to false.
           */
          public boolean completeUnlessQueued() {
              if (queuedType == State.UNQUEUED) {
                  running = false;
                  return false;
              } else {
                  return true;
              }
          }
      }
      private final InstanceState instanceState = new InstanceState();
  
      BookieInfoReader(BookKeeper bk,
                       ClientConfiguration conf,
                       ScheduledExecutorService scheduler) {
          this.bk = bk;
          this.conf = conf;
          this.scheduler = scheduler;
      }
  
      public void start() {
          this.bk
              .getMetadataClientDriver()
              .getRegistrationClient()
              .watchWritableBookies(bookies -> availableBookiesChanged(bookies.getValue()));
          scheduler.scheduleAtFixedRate(new Runnable() {
              @Override
              public void run() {
                  synchronized (BookieInfoReader.this) {
                      if (LOG.isDebugEnabled()) {
                          LOG.debug("Running periodic BookieInfo scan");
                      }
                      try {
                          Collection<BookieId> updatedBookies = bk.bookieWatcher.getBookies();
                          bookieInfoMap.updateBookies(updatedBookies);
                      } catch (BKException e) {
                          LOG.info("Got exception while querying bookies from watcher, rerunning after {}s",
                                   conf.getGetBookieInfoRetryIntervalSeconds(), e);
                          scheduler.schedule(this, conf.getGetBookieInfoRetryIntervalSeconds(), TimeUnit.SECONDS);
                          return;
                      }
                      if (instanceState.tryStartFull()) {
                          getReadWriteBookieInfo();
                      }
                  }
              }
          }, 0, conf.getGetBookieInfoIntervalSeconds(), TimeUnit.SECONDS);
      }
  
      private void submitTask() {
          scheduler.submit(() -> getReadWriteBookieInfo());
      }
  
      private void submitTaskWithDelay(int delaySeconds) {
          scheduler.schedule(() -> getReadWriteBookieInfo(), delaySeconds, TimeUnit.SECONDS);
      }
  
      synchronized void availableBookiesChanged(Set<BookieId> updatedBookiesList) {
          if (LOG.isInfoEnabled()) {
              LOG.info("Scheduling bookie info read due to changes in available bookies.");
          }
          bookieInfoMap.updateBookies(updatedBookiesList);
          if (instanceState.tryStartPartial()) {
              submitTask();
          }
      }
  
      /**
       * Method to allow tests to block until bookie info is available.
       *
       * @param bookie to lookup
       * @return None if absent, free disk space if present
       */
      synchronized Optional<Long> getFreeDiskSpace(BookieId bookie) {
          BookieInfo bookieInfo = bookieInfoMap.getInfo(bookie);
          if (bookieInfo != null) {
              return Optional.of(bookieInfo.getFreeDiskSpace());
          } else {
              return Optional.empty();
          }
      }
  
      /* State to track scan execution progress as callbacks come in */
      private int totalSent = 0;
      private int completedCnt = 0;
      private int errorCnt = 0;
  
      /**
       * Performs scan described by instanceState using the cached bookie information
       * in bookieInfoMap.
       */
      synchronized void getReadWriteBookieInfo() {
          State queuedType = instanceState.getAndClearQueuedType();
          Collection<BookieId> toScan;
          if (queuedType == State.FULL) {
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Doing full scan");
              }
              toScan = bookieInfoMap.getFullScanTargets();
          } else if (queuedType == State.PARTIAL) {
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Doing partial scan");
              }
              toScan = bookieInfoMap.getPartialScanTargets();
          } else {
              if (LOG.isErrorEnabled()) {
                  LOG.error("Invalid state, queuedType cannot be UNQUEUED in getReadWriteBookieInfo");
              }
              assert(queuedType != State.UNQUEUED);
              return;
          }
  
          BookieClient bkc = bk.getBookieClient();
          final long requested = BookkeeperProtocol.GetBookieInfoRequest.Flags.TOTAL_DISK_CAPACITY_VALUE
                                 | BookkeeperProtocol.GetBookieInfoRequest.Flags.FREE_DISK_SPACE_VALUE;
          totalSent = 0;
          completedCnt = 0;
          errorCnt = 0;
  
          if (LOG.isDebugEnabled()) {
              LOG.debug("Getting bookie info for: {}", toScan);
          }
          for (BookieId b : toScan) {
              bkc.getBookieInfo(b, requested,
                      new GetBookieInfoCallback() {
                          void processReadInfoComplete(int rc, BookieInfo bInfo, Object ctx) {
                              synchronized (BookieInfoReader.this) {
                                  BookieId b = (BookieId) ctx;
                                  if (rc != BKException.Code.OK) {
                                      if (LOG.isErrorEnabled()) {
                                          LOG.error("Reading bookie info from bookie {} failed due to {}",
                                                  b, BKException.codeLogger(rc));
                                      }
                                      // We reread bookies missing from the map each time, so remove to ensure
                                      // we get to it on the next scan
                                      bookieInfoMap.clearInfo(b);
                                      errorCnt++;
                                  } else {
                                      if (LOG.isDebugEnabled()) {
                                          LOG.debug("Bookie Info for bookie {} is {}", b, bInfo);
                                      }
                                      bookieInfoMap.gotInfo(b, bInfo);
                                  }
                                  completedCnt++;
                                  if (totalSent == completedCnt) {
                                      onExit();
                                  }
                              }
                          }
                          @Override
                          public void getBookieInfoComplete(final int rc, final BookieInfo bInfo, final Object ctx) {
                              scheduler.submit(
                                  new Runnable() {
                                      @Override
                                      public void run() {
                                          processReadInfoComplete(rc, bInfo, ctx);
                                      }
                                  });
                          }
                      }, b);
              totalSent++;
          }
          if (totalSent == 0) {
              onExit();
          }
      }
  
      void onExit() {
          bk.placementPolicy.updateBookieInfo(bookieInfoMap.getBookieMap());
          if (errorCnt > 0) {
              if (LOG.isInfoEnabled()) {
                  LOG.info("Rescheduling in {}s due to errors", conf.getGetBookieInfoIntervalSeconds());
              }
              instanceState.tryStartPartial();
              submitTaskWithDelay(conf.getGetBookieInfoRetryIntervalSeconds());
          } else if (instanceState.completeUnlessQueued()) {
              if (LOG.isInfoEnabled()) {
                  LOG.info("Rescheduling, another scan is pending");
              }
              submitTask();
          }
      }
  
      Map<BookieId, BookieInfo> getBookieInfo() throws BKException, InterruptedException {
          BookieClient bkc = bk.getBookieClient();
          final AtomicInteger totalSent = new AtomicInteger();
          final AtomicInteger totalCompleted = new AtomicInteger();
          final ConcurrentMap<BookieId, BookieInfo> map =
              new ConcurrentHashMap<BookieId, BookieInfo>();
          final CountDownLatch latch = new CountDownLatch(1);
          long requested = BookkeeperProtocol.GetBookieInfoRequest.Flags.TOTAL_DISK_CAPACITY_VALUE
                           | BookkeeperProtocol.GetBookieInfoRequest.Flags.FREE_DISK_SPACE_VALUE;
  
          Collection<BookieId> bookies;
          bookies = bk.bookieWatcher.getBookies();
          bookies.addAll(bk.bookieWatcher.getReadOnlyBookies());
          if (bookies.isEmpty()) {
              return map;
          }
          totalSent.set(bookies.size());
          for (BookieId b : bookies) {
              bkc.getBookieInfo(b, requested, new GetBookieInfoCallback() {
                          @Override
                          public void getBookieInfoComplete(int rc, BookieInfo bInfo, Object ctx) {
                              BookieId b = (BookieId) ctx;
                              if (rc != BKException.Code.OK) {
                                  if (LOG.isErrorEnabled()) {
                                      LOG.error("Reading bookie info from bookie {} failed due to {}",
                                              b, BKException.codeLogger(rc));
                                  }
                              } else {
                                  if (LOG.isDebugEnabled()) {
                                      LOG.debug("Free disk space on bookie {} is {}.", b, bInfo.getFreeDiskSpace());
                                  }
                                  map.put(b, bInfo);
                              }
                              if (totalCompleted.incrementAndGet() == totalSent.get()) {
                                  latch.countDown();
                              }
                          }
                      }, b);
          }
          try {
              latch.await();
          } catch (InterruptedException e) {
              Thread.currentThread().interrupt();
              LOG.error("Received InterruptedException ", e);
              throw e;
          }
          return map;
      }
  }
  