package com.atlassian.cache.ehcache;
   
   import java.io.Serializable;
   import java.util.Collection;
   import java.util.Map;
   import java.util.concurrent.locks.Lock;
   import java.util.concurrent.locks.ReadWriteLock;
   import java.util.concurrent.locks.ReentrantReadWriteLock;
   
  import javax.annotation.Nullable;
  
  import com.atlassian.cache.CacheLoader;
  
  import com.google.common.base.Function;
  import com.google.common.base.Throwables;
  import com.google.common.collect.Multimap;
  import com.google.common.collect.Multimaps;
  
  import net.sf.ehcache.CacheException;
  import net.sf.ehcache.Ehcache;
  import net.sf.ehcache.Element;
  import net.sf.ehcache.concurrent.LockType;
  import net.sf.ehcache.concurrent.Sync;
  import net.sf.ehcache.constructs.blocking.BlockingCache;
  import net.sf.ehcache.constructs.blocking.SelfPopulatingCache;
  
  import static com.atlassian.util.concurrent.Assertions.notNull;
  
  /**
   * A rewrite of {@link SelfPopulatingCache} with stronger concurrency guarantees.
   * Unfortunately, since the locks are local and cache replication happens at the Ehcache layer,
   * bypassing these locks, this is still not good enough for clusters.  It is, however,
   * an improvement.
   *
   * @param <K> the cache's key type
   * @param <V> the cache's value type
   */
  public class LoadingCache<K,V> extends BlockingCache
  {
      /**
       * Default value for number of mutexes in StripedReadWriteLockSync.DEFAULT_NUMBER_OF_MUTEXES is 2048,
       * however it's an overkill for Cloud where cache access is mostly request based and does not have
       * high concurrent usage. Reducing to 64 as default with system property override option.
       */
      private static final int DEFAULT_NUMBER_OF_MUTEXES = Integer.getInteger(LoadingCache.class.getName() + '.' + "DEFAULT_NUMBER_OF_MUTEXES", 64);
  
      private final CacheLoader<K, V> loader;
  
      // This needs to be fair to ensure that removeAll does not starve for a busy cache
      private final ReadWriteLock loadVsRemoveAllLock = new ReentrantReadWriteLock(true);
  
      public LoadingCache(final Ehcache cache, final CacheLoader<K, V> loader) throws CacheException
      {
          super(cache, DEFAULT_NUMBER_OF_MUTEXES);
          this.loader = notNull("loader", loader);
      }
  
      Lock loadLock()
      {
          return loadVsRemoveAllLock.readLock();
      }
  
      Lock removeAllLock()
      {
          return loadVsRemoveAllLock.writeLock();
      }
  
      @Override
      public Element get(final Object key)
      {
          if (key == null)
          {
              throw new NullPointerException("null keys are not permitted");
          }
          Element element = super.get(key);
          return (element != null) ? element : loadValueAndReleaseLock(key);
      }
  
      /**
       * Handle a cache miss by loading the value and releasing the write lock.
       * <p>
       * On a cache miss, {@link BlockingCache#get(Object) super.get(key)} returns {@code null} with the write
       * lock still held.  It is the caller's (our) responsibility to load the value and
       * {@link BlockingCache#put(Element) put} it into the cache, which will implicitly release the lock.
       * The lock must be released regardless of whether or not the load operation throws an exception.
       * </p>
       *
       * @param key the key for which we must load the value
       * @return a cache element representing the key and the corresponding value that was loaded for it
       */
      private Element loadValueAndReleaseLock(final Object key)
      {
          Element element;
          loadLock().lock();
          try
          {
              V value = null;
              try
              {
                 value = getFromLoader(key);
             }
             finally
             {
                 element = putLoadedValueAndReleaseLock(key, value);
             }
         }
         finally
         {
             // It isn't safe to let removeAll proceed until after we have done our put, so this
             // coarser read lock has to be held until the very end. :(
             loadLock().unlock();
         }
         return element;
     }
 
     @Nullable
     private Element putLoadedValueAndReleaseLock(Object key, V value)
     {
         // Putting the new element into the cache implicitly releases the lock for us
         if (value != null)
         {
             final Element element = new Element(key, value);
             put(element);
             return element;
         }
 
         // If the value is null, then loadValueAndReleaseLock is throwing an exception.  We still need to
         // release the lock, but the actual return value does not matter.  Note that unlike SelfPopulatingCache,
         // we are not using put(new Element(key, null)) to do this.  That would do an explicit removal (and with
         // replication!) which seems pretty pointless.
         final Sync lock = getLockForKey(key);
         if (lock.isHeldByCurrentThread(LockType.WRITE))
         {
             lock.unlock(LockType.WRITE);
         }
         return null;
     }
 
     @SuppressWarnings({ "ConstantConditions", "unchecked" })
     private V getFromLoader(Object key)
     {
         final V value;
         try
         {
             value = loader.load((K)key);
         }
         catch (final RuntimeException re)
         {
             put(new Element(key, null));
             throw propagate(key, re);
         }
         catch (final Error err)
         {
             put(new Element(key, null));
             throw propagate(key, err);
         }
 
         if (value == null)
         {
             throw new CacheException("CacheLoader returned null for key " + key);
         }
         return value;
     }
 
 
 
     // Make sure we acquire the write lock when removing a key.  BlockingCache should really be
     // doing this itself, but it isn't.  This prevents remove from overlapping with a lazy load
     // for the same key.
 
     @Override
     public boolean remove(Serializable key, boolean doNotNotifyCacheReplicators)
     {
         final Sync sync = getLockForKey(key);
         sync.lock(LockType.WRITE);
         try
         {
             return super.remove(key, doNotNotifyCacheReplicators);
         }
         finally
         {
             sync.unlock(LockType.WRITE);
         }
     }
 
     @Override
     public boolean remove(Serializable key)
     {
         final Sync sync = getLockForKey(key);
         sync.lock(LockType.WRITE);
         try
         {
             return super.remove(key);
         }
         finally
         {
             sync.unlock(LockType.WRITE);
         }
     }
 
     @Override
     public boolean remove(Object key)
     {
         final Sync sync = getLockForKey(key);
         sync.lock(LockType.WRITE);
         try
         {
             return super.remove(key);
         }
         finally
         {
             sync.unlock(LockType.WRITE);
         }
     }
 
     @Override
     public boolean remove(Object key, final boolean doNotNotifyCacheReplicators)
     {
         final Sync sync = getLockForKey(key);
         sync.lock(LockType.WRITE);
         try
         {
             return super.remove(key, doNotNotifyCacheReplicators);
         }
         finally
         {
             sync.unlock(LockType.WRITE);
         }
     }
 
 
 
     // removeAll(Collection) and removeAll(Collection, boolean) can specify multiple keys.  As with the
     // single-key remove methods, these need to acquire write locks to prevent them from overlapping with
     // lazy loads.  We could take the simple approach of iterating and delegating each one to remove(Object),
     // but it is easy enough to group them into lock stripes first, and this seems like a wise precaution
     // against a large cache removing several values at once (say, through expiration).
 
     @Override
     public void removeAll(Collection<?> keys)
     {
         removeGroupedBySync(keys, new RemoveCallback()
         {
             @Override
             public void removeUnderLock(Collection<?> keysForSync)
             {
                 underlyingCache.removeAll(keysForSync);
             }
         });
     }
 
     @Override
     public void removeAll(Collection<?> keys, final boolean doNotNotifyCacheReplicators)
     {
         removeGroupedBySync(keys, new RemoveCallback()
         {
             @Override
             public void removeUnderLock(Collection<?> keysForSync)
             {
                 underlyingCache.removeAll(keysForSync, doNotNotifyCacheReplicators);
             }
         });
     }
 
     /**
      * Partitions the supplied keys into subsets that share the same lock, then calls
      * {@link #removeGroupedBySync(Multimap, RemoveCallback)} with the results so that the callback
      * can be applied to each subset with the corresponding write lock held.
      *
      * @param allKeys the keys to be grouped into subsets that share a lock, then removed from the cache
      * @param callback the removal function to apply each lock's subset of keys while that write lock is held
      */
     private void removeGroupedBySync(Collection<?> allKeys, RemoveCallback callback)
     {
         final Multimap<Sync,?> map = Multimaps.index(allKeys, new Function<Object,Sync>()
         {
             @Override
             public Sync apply(Object key)
             {
                 return getLockForKey(key);
             }
         });
         removeGroupedBySync(map, callback);
     }
 
     /**
      * For each sync, acquires the write lock and calls {@link RemoveCallback#removeUnderLock(Collection)} with
      * the corresponding list of keys, then releases the lock.
      *
      * @param keysBySync the mapping of locks to each lock's corresponding keys
      * @param callback the removal function to apply each lock's subset of keys while that write lock is held
      */
     private static <K> void removeGroupedBySync(Multimap<Sync,K> keysBySync, RemoveCallback callback)
     {
         for (Map.Entry<Sync,Collection<K>> entry : keysBySync.asMap().entrySet())
         {
             final Sync sync = entry.getKey();
             final Collection<K> keysUsingThisSync = entry.getValue();
             sync.lock(LockType.WRITE);
             try
             {
                 callback.removeUnderLock(keysUsingThisSync);
             }
             finally
             {
                 sync.unlock(LockType.WRITE);
             }
         }
     }
 
 
 
     // The removeAll functions would have to acquire every single striped lock individually to prevent
     // overlap with lazy loaders, and this probably isn't acceptable.  Instead, we pay the cost of an
     // additional R/W lock that loaders can share and removeAll must acquire exclusively.
 
     @Override
     public void removeAll()
     {
         removeAllLock().lock();
         try
         {
             super.removeAll();
         }
         finally
         {
             removeAllLock().unlock();
         }
     }
 
     @Override
     public void removeAll(boolean doNotNotifyCacheReplicators)
     {
         removeAllLock().lock();
         try
         {
             super.removeAll(doNotNotifyCacheReplicators);
         }
         finally
         {
             removeAllLock().unlock();
         }
     }
 
 
 
     private static RuntimeException propagate(Object key, Throwable e)
     {
         Throwables.propagateIfInstanceOf(e, CacheException.class);
         throw new CacheException("Could not fetch object for cache entry with key \"" + key + "\".", e);
     }
 
     interface RemoveCallback
     {
         void removeUnderLock(Collection<?> keys);
     }
 }