output.txt

Stormpot-pool-1-thread-1 - Creating object for key: key1 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e57fd78 BLAZEPOOL
Stormpot-pool-1-thread-1 - Creating object for key: key1 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@37bd2157 BLAZEPOOL
Stormpot-pool-1-thread-2 - Creating object for key: key2 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e94fa46 QUEUEPOOL
Stormpot-pool-1-thread-2 - Creating object for key: key2 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@7fa9f07d QUEUEPOOL
Test worker - Invalidating (expire + release) com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e57fd78 BLAZEPOOL
Test worker - Invalidating (expire + release) com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e94fa46 QUEUEPOOL
Stormpot-pool-1-thread-1 - Destroying object for key: key1 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@37bd2157 BLAZEPOOL
Stormpot-pool-1-thread-1 - Closing object com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@37bd2157 BLAZEPOOL
Stormpot-pool-1-thread-2 - Destroying object for key: key2 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@7fa9f07d QUEUEPOOL
Stormpot-pool-1-thread-2 - Closing object com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@7fa9f07d QUEUEPOOL
Stormpot-pool-1-thread-2 - Destroying object for key: key2 com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e94fa46 QUEUEPOOL
Stormpot-pool-1-thread-2 - Closing object com.airbnb.oyster.common.thrift.StormPotSimpleTest$TestPoolObject@3e94fa46 QUEUEPOOL

StormPotSimpleTest.java

import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import org.junit.Before;
import org.junit.Test;
import stormpot.*;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

public class StormPotSimpleTest {
  private Map<String, TestPoolState> poolMap = new ConcurrentHashMap<>();
  enum PoolType {
    BLAZEPOOL,
    QUEUEPOOL
  }

  public static class TestPoolObject implements Poolable {
    private Slot slot;
    private TestPoolState state;

    public TestPoolObject(Slot slot, TestPoolState state) {
      this.slot = slot;
      this.state = state;
    }

    public void invalidate() {
      if (state.isShutdown()) {
        System.out.println(Thread.currentThread().getName() + " - Invalidating (just release) " + this + " " + state.type.toString());
        slot.release(this);
      } else {
        System.out.println(Thread.currentThread().getName() + " - Invalidating (expire + release) " + this + " " + state.type.toString());
        slot.expire(this);
        slot.release(this);
      }
    }

    @Override
    public void release() {
      System.out.println(Thread.currentThread().getName() + " - Releasing (just release) " + this  + " "  + state.type.toString());
      slot.release(this);
    }

    public void close() {
      System.out.println(Thread.currentThread().getName() + " - Closing object " + this + " " + state.type.toString());
      state = null;
      slot = null;
    }
  }

  private class TestPoolState {
    private BlazePool<TestPoolObject> blazePool = null;
    private QueuePool<TestPoolObject> queuePool = null;
    public PoolType type;

    private TestPoolState(String key, PoolType type) {
      Config config = new Config();
      config.setSize(2);
      config.setExpiration(new Expiration() {
        @Override
        public boolean hasExpired(SlotInfo info) throws Exception {
          return false;
        }
      });
      config.setBackgroundExpirationEnabled(false);
      config.setPreciseLeakDetectionEnabled(true);
      config.setAllocator(new TestAllocator(key));
      this.type = type;
      if (type == PoolType.BLAZEPOOL) {
        this.blazePool = new BlazePool<>(config);
      } else {
        this.queuePool = new QueuePool<>(config);
      }
    }

    private boolean isShutdown() {
      if (type == PoolType.BLAZEPOOL) {
        return blazePool == null || blazePool.isShutDown();
      } else {
        return queuePool == null || queuePool.isShutDown();
      }
    }

    private TestPoolObject claim(int timeoutMS)
        throws InterruptedException, PoolException {
      if (type == PoolType.BLAZEPOOL) {
        return blazePool.claim(new Timeout(timeoutMS, TimeUnit.MILLISECONDS));
      } else {
        return queuePool.claim(new Timeout(timeoutMS, TimeUnit.MILLISECONDS));
      }
    }


    private Completion close() throws Exception {
      Completion poolCompletion;
      if (!isShutdown()) {
        if (type == PoolType.BLAZEPOOL) {
          poolCompletion = blazePool.shutdown();
          blazePool = null;
        } else {
          poolCompletion = queuePool.shutdown();
          queuePool = null;
        }
        return poolCompletion;
      }
      return null;
    }
  }


  @RequiredArgsConstructor
  private class TestAllocator implements Allocator<TestPoolObject> {
    @NonNull private final String key;

    @Override
    public TestPoolObject allocate(Slot slot) throws Exception {
      TestPoolObject obj = new TestPoolObject(slot, poolMap.get(key));
      System.out.println(Thread.currentThread().getName() + " - Creating object for key: " + key + " " + obj + " " + poolMap.get(key).type.toString());
      return obj;
    }

    @Override
    public void deallocate(TestPoolObject poolObject) throws Exception {
      System.out.println(Thread.currentThread().getName() + " - Destroying object for key: " + key  + " " + poolObject + " " + poolObject.state.type.toString());
      poolObject.close();
    }
  }

  @Before
  public void setUp() throws Exception {
    // do not initialize pool
    poolMap.putIfAbsent("key1", new TestPoolState("key1", PoolType.BLAZEPOOL));
    poolMap.putIfAbsent("key2", new TestPoolState("key2", PoolType.QUEUEPOOL));
    Thread.sleep(100);
  }

  @Test
  public void testInvalidateReleaseShutdown() throws Exception {
    //TestPoolObject p11 = poolMap.get("key1").claim(10);
    TestPoolObject p12 = poolMap.get("key1").claim(10);

    //TestPoolObject p21 = poolMap.get("key2").claim(10);
    TestPoolObject p22 = poolMap.get("key2").claim(10);

    // invalidate p12, p22
    p12.invalidate();
    p22.invalidate();

    //Thread.sleep(1000);
    // shutdown pools
    Completion poolOneCompletion = poolMap.get("key1").close();
    poolMap.remove("key1");
    poolOneCompletion.await(new Timeout(1000, TimeUnit.MILLISECONDS));

    Completion poolTwoCompletion = poolMap.get("key2").close();
    poolMap.remove("key2");
    poolTwoCompletion.await(new Timeout(1000, TimeUnit.MILLISECONDS));

    // create new pool
    //TestPoolState oldState = poolMap.put("key1", new TestPoolState("key1"));
    //claim
    //p12 =  poolMap.get("key1").claim(30);
    //TestPoolObject p3 =  poolMap.get("key1").claim(30);

    // return p11, p21 to old pool
    //p11.release();
    //p21.release();

    Thread.sleep(10000);
  }
}