package com.atlassian.vcache.internal.core;
   
   import org.slf4j.Logger;
   import org.slf4j.LoggerFactory;
   
   import java.util.ArrayList;
   import java.util.Collection;
   import java.util.Collections;
   import java.util.function.BiFunction;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Helper class that guards against a thread recursively performing the same key. This is useful to guard
   * against behaviour like a thread calling {@link java.util.concurrent.ConcurrentHashMap#compute(Object, BiFunction)}
   * recursively.
   *
   * @param <T> the key type.
   * @since 1.8.2
   */
  public class RecursionDetector<T> {
      private static final Logger log = LoggerFactory.getLogger(RecursionDetector.class);
  
      /**
       * Used to track inflight calls.
       */
      private final ThreadLocal<Collection<T>> INFLIGHT = new ThreadLocal<Collection<T>>() {
          @Override
          protected Collection<T> initialValue() {
              // Using an ArrayList as I believe that size of the collection will be very small, and hence the
              // performance will be faster than a Set.
              return new ArrayList<>();
          }
      };
  
      public Guard<T> guardOn(T key) {
          return new Guard<>(Collections.singletonList(key), INFLIGHT.get());
      }
  
      public Guard<T> guardOn(Iterable<T> keys) {
          return new Guard<>(keys, INFLIGHT.get());
      }
  
      /**
       * Helper class to detect recursive calls to the same key.
       *
       * @param <T> the key type.
       */
      public static class Guard<T> implements AutoCloseable {
          private final Iterable<T> keys;
          private final Collection<T> active;
  
          private Guard(Iterable<T> keys, Collection<T> active) {
              this.keys = requireNonNull(keys);
              this.active = requireNonNull(active);
  
              // first confirm none of the keys are being used, otherwise could leave the threadlocal in a bad state.
              // Naive implementation would loop once, checking each key and then adding them. If an error was detected,
              // then we would need to remove all the added keys, before throwing the exception.
              keys.forEach(k -> {
                  if (active.contains(k)) {
                      log.error("Detected recursive call for key {}.", k);
                      throw new IllegalStateException("Recursive call with key: " + k);
                  }
              });
  
              // now add all the keys
              keys.forEach(active::add);
          }
  
          @Override
          public void close() {
              keys.forEach(active::remove);
          }
      }
  }