package com.atlassian.cache.memory;
   
   import java.util.concurrent.ConcurrentHashMap;
   import java.util.concurrent.ConcurrentMap;
   import java.util.concurrent.locks.Lock;
   import java.util.concurrent.locks.ReadWriteLock;
   import java.util.concurrent.locks.ReentrantReadWriteLock;
   
   import javax.annotation.Nonnull;
  
  import com.atlassian.cache.impl.OneShotLatch;
  
  import com.google.common.cache.CacheLoader;
  
  import static com.atlassian.util.concurrent.Assertions.notNull;
  
  /**
   *
   * @since v2.4.5
   */
  public class BlockingCacheLoader<K,V> extends CacheLoader<K,V>
  {
      private final ConcurrentMap<K, OneShotLatch> barriers = new ConcurrentHashMap<K, OneShotLatch>(16);
      private final CacheLoader<K, V> delegate;
  
      // This needs to be fair to ensure that removeAll does not starve for a busy cache
      private final ReadWriteLock loadVsRemoveAllLock = new ReentrantReadWriteLock(true);
  
      BlockingCacheLoader(final CacheLoader<K, V> delegate)
      {
          this.delegate = notNull("delegate", delegate);
      }
  
      Lock loadLock()
      {
          return loadVsRemoveAllLock.readLock();
      }
  
      Lock removeAllLock()
      {
          return loadVsRemoveAllLock.writeLock();
      }
  
      @SuppressWarnings("LockAcquiredButNotSafelyReleased")
      @Override
      public V load(@Nonnull K key) throws Exception
      {
          // to be released by postGetCleanup
          acquire(key);
          loadLock().lock();
          return delegate.load(key);
      }
  
      void postGetCleanup(@Nonnull K key)
      {
          final OneShotLatch barrier = barriers.get(key);
          if (barrier != null && barrier.isHeldByCurrentThread())
          {
              // At this point, we know that Guava's LoadingValueReference has been replaced and we no longer
              // need to hold up removeAll, so release the coarser lock first.
              loadLock().unlock();
              barriers.remove(key);
              barrier.release();
          }
      }
  
      OneShotLatch acquire(@Nonnull K key)
      {
          final OneShotLatch barrier = new OneShotLatch();
          while (true)
          {
              final OneShotLatch existing = barriers.putIfAbsent(key, barrier);
  
              // successfully raised a new barrier
              if (existing == null)
              {
                  return barrier;
              }
  
              // There is an existing barrier that happens-before us; wait for it to be released.
              // There is no need to attempt a remove or replace to evict it from the barriers map;
              // The thread that owned it would have taken care of that before releasing it.
              existing.await();
          }
      }
  
      void release(@Nonnull K key)
      {
          final OneShotLatch existing = barriers.get(key);
          if (existing == null || !existing.isHeldByCurrentThread())
          {
              throw new IllegalMonitorStateException("existing=" + existing);
          }
          barriers.remove(key);
          existing.release();
      }
  
  
  }