package com.atlassian.vcache.internal.redis;
   
   import com.atlassian.annotations.tenancy.Tenantless;
   import com.atlassian.fugue.Pair;
   import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
   import org.slf4j.Logger;
   import org.slf4j.LoggerFactory;
   import redis.clients.jedis.Jedis;
   import redis.clients.jedis.JedisPool;
  
  import java.net.URI;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.Consumer;
  
  import static com.atlassian.fugue.Pair.pair;
  
  /**
   * A specialised {@link JedisPool} which binds the currently in-use {@link Jedis} instance to the requesting thread,
   * so that repeated calls to {@link #getResource()} without a corresponding {@link #returnResource(Jedis)} will
   * return the same instance. A reference count is kept for each thread-bound instance, and when an instance is returned
   * and the reference count reaches zero, the instance is returned to the pool for real.
   *
   * @since 1.10.1
   */
  public class ThreadAffinityJedisPool extends JedisPool {
      private static final Logger log = LoggerFactory.getLogger(ThreadAffinityJedisPool.class);
  
      @Tenantless(reason = "Thread-local data is used only to coordinate the sharing of a Jedis client over the short-term life of a VCache operation.")
      private final ThreadLocal<Pair<Jedis, AtomicInteger>> threadBoundResources = new ThreadLocal<>();
  
      public ThreadAffinityJedisPool() {
          super();
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host) {
          super(poolConfig, host);
      }
  
      public ThreadAffinityJedisPool(String host, int port) {
          super(host, port);
      }
  
      public ThreadAffinityJedisPool(String host) {
          super(host);
      }
  
      public ThreadAffinityJedisPool(URI uri) {
          super(uri);
      }
  
      public ThreadAffinityJedisPool(URI uri, int timeout) {
          super(uri, timeout);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port, int timeout, String password) {
          super(poolConfig, host, port, timeout, password);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port) {
          super(poolConfig, host, port);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port, int timeout) {
          super(poolConfig, host, port, timeout);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port, int timeout, String password, int database) {
          super(poolConfig, host, port, timeout, password, database);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port, int timeout, String password, int database, String clientName) {
          super(poolConfig, host, port, timeout, password, database, clientName);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, String host, int port, int connectionTimeout, int soTimeout, String password, int database, String clientName) {
          super(poolConfig, host, port, connectionTimeout, soTimeout, password, database, clientName);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, URI uri) {
          super(poolConfig, uri);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, URI uri, int timeout) {
          super(poolConfig, uri, timeout);
      }
  
      public ThreadAffinityJedisPool(GenericObjectPoolConfig poolConfig, URI uri, int connectionTimeout, int soTimeout) {
          super(poolConfig, uri, connectionTimeout, soTimeout);
      }
  
      @Override
      public Jedis getResource() {
          final Jedis jedis = getOrCreateThreadBound().left();
          final int newCount = getOrCreateThreadBound().right().incrementAndGet();
          log.debug("Using thread-bound {}, reference count is now {}", jedis, newCount);
          return jedis;
      }
  
      private Pair<Jedis, AtomicInteger> getOrCreateThreadBound() {
         final Pair<Jedis, AtomicInteger> threadBoundResource = threadBoundResources.get();
         if (threadBoundResource == null) {
             log.debug("No thread-bound Jedis instance found, fetching one from the pool");
             final Jedis jedis = super.getResource();
             threadBoundResources.set(pair(jedis, new AtomicInteger(0)));
             log.debug("Bound {} to thread with reference count 0", jedis);
         }
         return threadBoundResources.get();
     }
 
     @Override
     public void returnBrokenResource(Jedis jedis) {
         returnInternal(jedis, super::returnBrokenResource);
     }
 
     @Override
     public void returnResource(Jedis jedis) {
         returnInternal(jedis, super::returnResource);
     }
 
     private void returnInternal(Jedis jedis, Consumer<Jedis> returnAction) {
         log.debug("Returning {}", jedis);
         if (threadBoundResources.get() == null) {
             throw new IllegalStateException("No thread-bound resource present");
         }
         if (jedis != threadBoundResources.get().left()) {
             throw new IllegalStateException("Thread-bound Jedis object is not the same as the object being returned");
         }
         final int newCount = threadBoundResources.get().right().decrementAndGet();
         if (newCount == 0) {
             log.debug("Reference count of thread-bound {} is now 0, returning to the pool", jedis);
             threadBoundResources.remove();
             returnAction.accept(jedis);
         } else {
             log.debug("Reference count of thread-bound {} is now {}, not returning to pool", jedis, newCount);
         }
     }
 }