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.kafka.server.purgatory;
  
  import org.apache.kafka.server.metrics.KafkaMetricsGroup;
  import org.apache.kafka.server.util.ShutdownableThread;
  import org.apache.kafka.server.util.timer.SystemTimer;
  import org.apache.kafka.server.util.timer.Timer;
  import org.apache.kafka.server.util.timer.TimerTask;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.locks.ReentrantLock;
  
  public class DelayedOperationPurgatory<T extends DelayedOperation> {
  
      private static final Logger LOG = LoggerFactory.getLogger(DelayedOperationPurgatory.class);
      private static final int SHARDS = 512; // Shard the watcher list to reduce lock contention
  
      private final KafkaMetricsGroup metricsGroup = new KafkaMetricsGroup("kafka.server", "DelayedOperationPurgatory");
      private final Map<String, String> metricsTags;
      private final List<WatcherList> watcherLists;
      // the number of estimated total operations in the purgatory
      private final AtomicInteger estimatedTotalOperations = new AtomicInteger(0);
      /* background thread expiring operations that have timed out */
      private final ExpiredOperationReaper expirationReaper;
      private final String purgatoryName;
      private final Timer timeoutTimer;
      private final int brokerId;
      private final int purgeInterval;
      private final boolean reaperEnabled;
      private final boolean timerEnabled;
  
      public DelayedOperationPurgatory(String purgatoryName, Timer timer, int brokerId, boolean reaperEnabled) {
          this(purgatoryName, timer, brokerId, 1000, reaperEnabled, true);
      }
  
      public DelayedOperationPurgatory(String purgatoryName, Timer timer, int brokerId, int purgeInterval) {
          this(purgatoryName, timer, brokerId, purgeInterval, true, true);
      }
  
      public DelayedOperationPurgatory(String purgatoryName, int brokerId) {
          this(purgatoryName, brokerId, 1000);
      }
  
      public DelayedOperationPurgatory(String purgatoryName, int brokerId, int purgeInterval) {
          this(purgatoryName, new SystemTimer(purgatoryName), brokerId, purgeInterval, true, true);
      }
  
      /**
       * A helper purgatory class for bookkeeping delayed operations with a timeout, and expiring timed out operations.
       */
      @SuppressWarnings("this-escape")
      public DelayedOperationPurgatory(String purgatoryName,
                                       Timer timeoutTimer,
                                       int brokerId,
                                       int purgeInterval,
                                       boolean reaperEnabled,
                                       boolean timerEnabled) {
          this.purgatoryName = purgatoryName;
          this.timeoutTimer = timeoutTimer;
          this.brokerId = brokerId;
          this.purgeInterval = purgeInterval;
          this.reaperEnabled = reaperEnabled;
          this.timerEnabled = timerEnabled;
          // The initialization of the expiration reaper thread is done after the brokerId and purgatoryName
          // are set, so that the thread name can include these values.
          this.expirationReaper = new ExpiredOperationReaper();
  
          watcherLists = new ArrayList<>(SHARDS);
          for (int i = 0; i < SHARDS; i++) {
              watcherLists.add(new WatcherList());
          }
          metricsTags = Map.of("delayedOperation", purgatoryName);
          metricsGroup.newGauge("PurgatorySize", this::watched, metricsTags);
          metricsGroup.newGauge("NumDelayedOperations", this::numDelayed, metricsTags);
          if (reaperEnabled) {
              expirationReaper.start();
          }
      }
  
      private WatcherList watcherList(DelayedOperationKey key) {
          return watcherLists.get(Math.abs(key.hashCode() % watcherLists.size()));
      }
  
      /**
       * Check if the operation can be completed, if not watch it based on the given watch keys
       * <br/>
       * Note that a delayed operation can be watched on multiple keys. It is possible that
       * an operation is completed after it has been added to the watch list for some, but
       * not all the keys. In this case, the operation is considered completed and won't
       * be added to the watch list of the remaining keys. The expiration reaper thread will
       * remove this operation from any watcher list in which the operation exists.
       *
       * @param operation the delayed operation to be checked
       * @param watchKeys keys for bookkeeping the operation
       * @return true iff the delayed operations can be completed by the caller
       */
      public <K extends DelayedOperationKey> boolean tryCompleteElseWatch(T operation, List<K> watchKeys) {
          if (watchKeys.isEmpty()) {
              throw new IllegalArgumentException("The watch key list can't be empty");
          }
  
          // The cost of tryComplete() is typically proportional to the number of keys. Calling tryComplete() for each key is
          // going to be expensive if there are many keys. Instead, we do the check in the following way through safeTryCompleteOrElse().
          // If the operation is not completed, we just add the operation to all keys. Then we call tryComplete() again. At
          // this time, if the operation is still not completed, we are guaranteed that it won't miss any future triggering
          // event since the operation is already on the watcher list for all keys.
          //
          // ==============[story about lock]==============
          // Through safeTryCompleteOrElse(), we hold the operation's lock while adding the operation to watch list and doing
          // the tryComplete() check. This is to avoid a potential deadlock between the callers to tryCompleteElseWatch() and
          // checkAndComplete(). For example, the following deadlock can happen if the lock is only held for the final tryComplete()
          // 1) thread_a holds readlock of stateLock from TransactionStateManager
          // 2) thread_a is executing tryCompleteElseWatch()
          // 3) thread_a adds op to watch list
          // 4) thread_b requires writelock of stateLock from TransactionStateManager (blocked by thread_a)
          // 5) thread_c calls checkAndComplete() and holds lock of op
          // 6) thread_c is waiting readlock of stateLock to complete op (blocked by thread_b)
          // 7) thread_a is waiting lock of op to call the final tryComplete() (blocked by thread_c)
          //
          // Note that even with the current approach, deadlocks could still be introduced. For example,
          // 1) thread_a calls tryCompleteElseWatch() and gets lock of op
          // 2) thread_a adds op to watch list
          // 3) thread_a calls op#tryComplete and tries to require lock_b
          // 4) thread_b holds lock_b and calls checkAndComplete()
          // 5) thread_b sees op from watch list
          // 6) thread_b needs lock of op
          // To avoid the above scenario, we recommend DelayedOperationPurgatory.checkAndComplete() be called without holding
          // any exclusive lock. Since DelayedOperationPurgatory.checkAndComplete() completes delayed operations asynchronously,
          // holding an exclusive lock to make the call is often unnecessary.
          if (operation.safeTryCompleteOrElse(() -> {
              watchKeys.forEach(key -> {
                  if (!operation.isCompleted())
                      watchForOperation(key, operation);
              });
              if (!watchKeys.isEmpty())
                  estimatedTotalOperations.incrementAndGet();
          })) {
              return true;
          }
  
          // if it cannot be completed by now and hence is watched, add to the timeout queue also
          if (!operation.isCompleted()) {
              if (timerEnabled)
                  timeoutTimer.add(operation);
              if (operation.isCompleted()) {
                  // cancel the timer task
                  operation.cancel();
              }
          }
          return false;
      }
  
      /**
       * Check if some delayed operations can be completed with the given watch key,
       * and if yes complete them.
       *
       * @return the number of completed operations during this process
       */
      public <K extends DelayedOperationKey> int checkAndComplete(K key) {
          WatcherList wl = watcherList(key);
          Watchers watchers;
          wl.watchersLock.lock();
          try {
              watchers = wl.watchersByKey.get(key);
          } finally {
              wl.watchersLock.unlock();
          }
          int numCompleted = watchers == null ? 0 : watchers.tryCompleteWatched();
  
          if (numCompleted > 0) {
              LOG.debug("Request key {} unblocked {} {} operations", key, numCompleted, purgatoryName);
          }
          return numCompleted;
      }
  
      /**
       * Return the total size of watch lists the purgatory. Since an operation may be watched
       * on multiple lists, and some of its watched entries may still be in the watch lists
       * even when it has been completed, this number may be larger than the number of real operations watched
       */
      public int watched() {
          int sum = 0;
          for (WatcherList watcherList : watcherLists) {
              sum += watcherList.allWatchers().stream().mapToInt(Watchers::countWatched).sum();
          }
          return sum;
      }
  
      /**
       * Return the number of delayed operations in the expiry queue
       */
      public int numDelayed() {
          return timeoutTimer.size();
      }
  
      /**
       * Cancel watching on any delayed operations for the given key. Note the operation will not be completed
       */
      public List<T> cancelForKey(DelayedOperationKey key) {
          WatcherList wl = watcherList(key);
          wl.watchersLock.lock();
          try {
              Watchers watchers = wl.watchersByKey.remove(key);
              if (watchers != null)
                  return watchers.cancel();
              else
                  return List.of();
          } finally {
              wl.watchersLock.unlock();
          }
      }
  
      /*
       * Watch the operation
       */
      private void watchForOperation(DelayedOperationKey key, T operation) {
          WatcherList wl = watcherList(key);
          wl.watchersLock.lock();
          try {
              Watchers watcher = wl.watchersByKey.computeIfAbsent(key, Watchers::new);
              watcher.watch(operation);
          } finally {
              wl.watchersLock.unlock();
          }
      }
  
      /*
       * Remove the key from watcher lists if its list is empty
       */
      private void removeKeyIfEmpty(DelayedOperationKey key, Watchers watchers) {
          WatcherList wl = watcherList(key);
          wl.watchersLock.lock();
          try {
              // if the current key is no longer correlated to the watchers to remove, skip
              if (wl.watchersByKey.get(key) != watchers)
                  return;
  
              if (watchers != null && watchers.isEmpty()) {
                  wl.watchersByKey.remove(key);
              }
          } finally {
              wl.watchersLock.unlock();
          }
      }
  
      /**
       * Shutdown the expiration reaper thread
       */
      public void shutdown() throws Exception {
          if (reaperEnabled) {
              expirationReaper.initiateShutdown();
              // improve shutdown time by waking up any ShutdownableThread blocked on poll by sending a no-op
              timeoutTimer.add(new TimerTask(0) {
                  @Override
                  public void run() {}
              });
              expirationReaper.awaitShutdown();
          }
          timeoutTimer.close();
          metricsGroup.removeMetric("PurgatorySize", metricsTags);
          metricsGroup.removeMetric("NumDelayedOperations", metricsTags);
      }
  
      /**
       * A list of operation watching keys
       */
      private class WatcherList {
          private final ConcurrentHashMap<DelayedOperationKey, Watchers> watchersByKey = new ConcurrentHashMap<>();
  
          private final ReentrantLock watchersLock = new ReentrantLock();
  
          /*
           * Return all the current watcher lists,
           * note that the returned watchers may be removed from the list by other threads
           */
          Collection<Watchers> allWatchers() {
              return watchersByKey.values();
          }
      }
  
      /**
       * A linked list of watched delayed operations based on some key
       */
      private class Watchers {
  
          private final ConcurrentLinkedQueue<T> operations = new ConcurrentLinkedQueue<>();
  
          private final DelayedOperationKey key;
          Watchers(DelayedOperationKey key) {
              this.key = key;
          }
  
          // count the current number of watched operations. This is O(n), so use isEmpty() if possible
          int countWatched() {
              return operations.size();
          }
  
          boolean isEmpty() {
              return operations.isEmpty();
          }
  
          // add the element to watch
          void watch(T t) {
              operations.add(t);
          }
  
          // traverse the list and try to complete some watched elements
          int tryCompleteWatched() {
              int completed = 0;
  
              Iterator<T> iter = operations.iterator();
              while (iter.hasNext()) {
                  T curr = iter.next();
                  if (curr.isCompleted()) {
                      // another thread has completed this operation, just remove it
                      iter.remove();
                  } else if (curr.safeTryComplete()) {
                      iter.remove();
                      completed += 1;
                  }
              }
  
              if (operations.isEmpty())
                  removeKeyIfEmpty(key, this);
  
              return completed;
          }
  
          List<T> cancel() {
              Iterator<T> iter = operations.iterator();
              List<T> cancelled = new ArrayList<>();
              while (iter.hasNext()) {
                  T curr = iter.next();
                  curr.cancel();
                  iter.remove();
                  cancelled.add(curr);
              }
              return cancelled;
          }
  
          // traverse the list and purge elements that are already completed by others
          int purgeCompleted() {
              int purged = 0;
  
              Iterator<T> iter = operations.iterator();
              while (iter.hasNext()) {
                  T curr = iter.next();
                  if (curr.isCompleted()) {
                      iter.remove();
                      purged += 1;
                  }
              }
  
              if (operations.isEmpty())
                  removeKeyIfEmpty(key, this);
  
              return purged;
          }
      }
  
      private void advanceClock(long timeoutMs) throws InterruptedException {
          timeoutTimer.advanceClock(timeoutMs);
  
          // Trigger a purge if the number of completed but still being watched operations is larger than
          // the purge threshold. That number is computed by the difference btw the estimated total number of
          // operations and the number of pending delayed operations.
          if (estimatedTotalOperations.get() - numDelayed() > purgeInterval) {
              // now set estimatedTotalOperations to delayed (the number of pending operations) since we are going to
              // clean up watchers. Note that, if more operations are completed during the cleanup, we may end up with
              // a little overestimated total number of operations.
              estimatedTotalOperations.getAndSet(numDelayed());
              LOG.debug("Begin purging watch lists");
              int purged = 0;
              for (WatcherList watcherList : watcherLists) {
                  purged += watcherList.allWatchers().stream().mapToInt(Watchers::purgeCompleted).sum();
              }
              LOG.debug("Purged {} elements from watch lists.", purged);
          }
      }
  
      /**
       * A background reaper to expire delayed operations that have timed out
       */
      private class ExpiredOperationReaper extends ShutdownableThread {
          ExpiredOperationReaper() {
              super("ExpirationReaper-" + brokerId + "-" + purgatoryName, false);
          }
  
          @Override
          public void doWork() {
              try {
                  advanceClock(200L);
              } catch (InterruptedException ie) {
                  throw new RuntimeException(ie);
              }
          }
      }
  }
  