/*
 * Copyright 2018 Netflix, Inc.
 *
 * Licensed 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 com.netflix.spinnaker.kork.jedis.telemetry;

import com.netflix.spectator.api.Registry;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

public class InstrumentedJedisPool extends JedisPool {

  private final Registry registry;
  private final JedisPool delegated;
  private final String poolName;

  private GenericObjectPool<Jedis> delegateInternalPool;

  public InstrumentedJedisPool(Registry registry, JedisPool delegated) {
    this(registry, delegated, "unnamed");
  }

  public InstrumentedJedisPool(Registry registry, JedisPool delegated, String poolName) {
    this.registry = registry;
    this.delegated = delegated;
    this.poolName = poolName;

    getInternalPoolReference();
  }

  @SuppressWarnings("unchecked")
  private void getInternalPoolReference() {
    try {
      Field f = delegated.getClass().getDeclaredField("internalPool");
      f.setAccessible(true);
      delegateInternalPool = (GenericObjectPool<Jedis>) f.get(delegated);
    } catch (NoSuchFieldException|IllegalAccessException e) {
      throw new IllegalStateException("Could not get reference to delegate's internal pool", e);
    }
  }

  @Override
  public Jedis getResource() {
    return new InstrumentedJedis(registry, delegated.getResource(), poolName);
  }

  @Override
  public void returnResourceObject(Jedis resource) {
    super.returnResourceObject(unwrapResource(resource));
  }

  @Override
  protected void returnBrokenResourceObject(Jedis resource) {
    super.returnBrokenResourceObject(unwrapResource(resource));
  }

  @Override
  public void close() {
    delegated.close();
  }

  @Override
  public boolean isClosed() {
    return delegated.isClosed();
  }

  @Override
  public void initPool(GenericObjectPoolConfig poolConfig, PooledObjectFactory<Jedis> factory) {
    delegated.initPool(poolConfig, factory);
  }

  @Override
  public void destroy() {
    delegated.destroy();
  }

  @Override
  protected void closeInternalPool() {
    try {
      Method m = delegated.getClass().getDeclaredMethod("closeInternalPool");
      m.setAccessible(true);
      m.invoke(delegated);
    } catch (NoSuchMethodException|InvocationTargetException|IllegalAccessException e) {
      throw new IllegalStateException("Could not invoke close internal pool on delegate", e);
    }
  }

  @Override
  public int getNumActive() {
    return delegateInternalPool.getNumActive();
  }

  @Override
  public int getNumIdle() {
    return delegateInternalPool.getNumIdle();
  }

  @Override
  public int getNumWaiters() {
    return delegateInternalPool.getNumWaiters();
  }

  @Override
  public long getMeanBorrowWaitTimeMillis() {
    return delegateInternalPool.getMeanBorrowWaitTimeMillis();
  }

  @Override
  public long getMaxBorrowWaitTimeMillis() {
    return delegateInternalPool.getMaxBorrowWaitTimeMillis();
  }

  @Override
  public void addObjects(int count) {
    delegated.addObjects(count);
  }

  private Jedis unwrapResource(Jedis jedis) {
    if (jedis instanceof InstrumentedJedis) {
      return ((InstrumentedJedis) jedis).unwrap();
    }
    return jedis;
  }
}

@@ -964,93 +964,8 @@
 eld;
-%0Aimport java.lang.reflect.InvocationTargetException;%0Aimport java.lang.reflect.Method;
 %0A%0Apu
@@ -1487,41 +1487,8 @@
 ame;
-%0A%0A    getInternalPoolReference();
 %0A  %7D
@@ -1532,20 +1532,40 @@
 private 
-void
+GenericObjectPool%3CJedis%3E
  getInte
@@ -1582,24 +1582,66 @@
 ference() %7B%0A
+    if (delegateInternalPool == null) %7B%0A  
     try %7B%0A  
@@ -1644,16 +1644,18 @@
 %7B%0A      
+ 
+ 
 Field f 
@@ -1717,16 +1717,18 @@
 ;%0A      
+  
 f.setAcc
@@ -1744,24 +1744,26 @@
 rue);%0A      
+ 
+ 
 delegateInte
@@ -1810,32 +1810,34 @@
 get(delegated);%0A
+  
     %7D catch (NoS
@@ -1880,32 +1880,34 @@
 tion e) %7B%0A      
+  
 throw new Illega
@@ -1977,33 +1977,74 @@
  pool%22, e);%0A    
-%7D
+  %7D%0A    %7D%0A    return delegateInternalPool;
 %0A  %7D%0A%0A  @Overrid
@@ -2692,48 +2692,98 @@
     
-delegated.initPool(poolConfig, factory);
+// Explicitly not initializing the pool here, as the delegated pool will initialize itself
 %0A  %7D
@@ -2911,336 +2911,99 @@
     
-try %7B%0A      Method m = delegated.getClass().getDeclaredMethod(%22closeInternalPool%22);%0A      m.setAccessible(true);%0A      m.invoke(delegated);%0A    %7D catch (NoSuchMethodException%7CInvocationTargetException%7CIllegalAccessException e) %7B%0A      throw new IllegalStateException(%22Could not invoke close internal pool on delegate%22, e);%0A    %7D
+// Explicitly not calling this; destroy and initPool are the only references to this method
 %0A  %7D
@@ -3053,32 +3053,27 @@
 %0A    return 
-delegate
+get
 InternalPool
@@ -3068,24 +3068,35 @@
 InternalPool
+Reference()
 .getNumActiv
@@ -3152,32 +3152,27 @@
 %0A    return 
-delegate
+get
 InternalPool
@@ -3167,24 +3167,35 @@
 InternalPool
+Reference()
 .getNumIdle(
@@ -3252,32 +3252,27 @@
 %0A    return 
-delegate
+get
 InternalPool
@@ -3267,24 +3267,35 @@
 InternalPool
+Reference()
 .getNumWaite
@@ -3370,32 +3370,27 @@
 %0A    return 
-delegate
+get
 InternalPool
@@ -3385,24 +3385,35 @@
 InternalPool
+Reference()
 .getMeanBorr
@@ -3501,32 +3501,27 @@
 %0A    return 
-delegate
+get
 InternalPool
@@ -3520,16 +3520,27 @@
 rnalPool
+Reference()
 .getMaxB

Fix NPE on delegate internal pool (#189)