From 778ac737e59342a22a028b3f500ee17dd0922816 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sun, 23 Dec 2018 16:40:06 -0600
Subject: [PATCH] Add *Pool.Remove(res)

---
 README.md    |   1 -
 pool.go      |  32 +++++++++++
 pool_test.go | 155 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 187 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index cec9981..278eda8 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,6 @@ Puddle is a generic resource pool library for Go.
 * Max resource idle time
 * Resource keep alive
 * Resource health check - keep alive and health check might be same thing
-* Remove resource (instead of just return to pool)
 * Reset pool
 * Shrink pool
 * Stress test
diff --git a/pool.go b/pool.go
index 3610f3b..f14f645 100644
--- a/pool.go
+++ b/pool.go
@@ -305,3 +305,35 @@ func (p *Pool) Return(res interface{}) {
 	p.cond.L.Unlock()
 	p.cond.Signal()
 }
+
+// Remove removes res from the pool and closes it. If res is not part of the
+// pool Remove will panic.
+func (p *Pool) Remove(res interface{}) {
+	p.cond.L.Lock()
+	defer p.cond.L.Unlock()
+
+	rw, present := p.allResources[res]
+	if !present {
+		panic("Remove called on resource that does not belong to pool")
+	}
+
+	delete(p.allResources, rw.resource)
+
+	// close the resource in the background
+	go func() {
+		err := p.closeRes(res)
+		if err != nil {
+			p.cond.L.Lock()
+			p.backgroundErrorHandler(err)
+			p.cond.L.Unlock()
+		}
+	}()
+
+	// Maintain min pool size (unless pool is already closed)
+	if !p.closed {
+		for len(p.allResources) < p.minSize {
+			createResChan, createErrChan := p.startCreate()
+			p.backgroundFinishCreate(createResChan, createErrChan)
+		}
+	}
+}
diff --git a/pool_test.go b/pool_test.go
index df7cb9a..96aeea4 100644
--- a/pool_test.go
+++ b/pool_test.go
@@ -223,6 +223,161 @@ func TestPoolReturnClosesResourcePoolIsAlreadyClosed(t *testing.T) {
 	assert.Equal(t, len(resources), closeCalls.Value())
 }
 
+func TestPoolRemovePanicsIfResourceNotPartOfPool(t *testing.T) {
+	var createCalls Counter
+	createFunc := func() (interface{}, error) {
+		return createCalls.Next(), nil
+	}
+	pool := puddle.NewPool(createFunc, stubCloseRes)
+
+	assert.Panics(t, func() { pool.Remove(42) })
+}
+
+func TestPoolRemoveRemovesResourceFromPool(t *testing.T) {
+	var createCalls Counter
+	createFunc := func() (interface{}, error) {
+		return createCalls.Next(), nil
+	}
+	pool := puddle.NewPool(createFunc, stubCloseRes)
+
+	res, err := pool.Get(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, 1, res)
+
+	assert.Equal(t, 1, pool.Size())
+	pool.Remove(res)
+	assert.Equal(t, 0, pool.Size())
+}
+
+func TestPoolRemoveRemovesResourceFromPoolAndStartsNewCreationToMaintainMinSize(t *testing.T) {
+	createCallsChan := make(chan int, 4)
+	closeCallsChan := make(chan int, 4)
+
+	waitForRead := func(ch chan int) bool {
+		select {
+		case <-ch:
+			return true
+		case <-time.NewTimer(time.Second).C:
+			return false
+		}
+	}
+
+	var createCalls Counter
+	createFunc := func() (interface{}, error) {
+		n := createCalls.Next()
+		createCallsChan <- n
+		return n, nil
+	}
+
+	var closeCalls Counter
+	closeFunc := func(interface{}) error {
+		n := closeCalls.Next()
+		closeCallsChan <- n
+		return nil
+	}
+
+	pool := puddle.NewPool(createFunc, closeFunc)
+
+	// Ensure there are 2 resources available in pool
+	{
+		r1, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		r2, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		pool.Return(r1)
+		pool.Return(r2)
+	}
+
+	assert.Equal(t, 2, pool.Size())
+	pool.SetMinSize(2)
+	assert.Equal(t, 2, pool.Size())
+
+	{
+		r1, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		r2, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		pool.Remove(r1)
+		pool.Remove(r2)
+	}
+
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(closeCallsChan))
+	require.True(t, waitForRead(closeCallsChan))
+
+	assert.Equal(t, 2, pool.Size())
+	assert.Equal(t, 4, createCalls.Value())
+	assert.Equal(t, 2, closeCalls.Value())
+}
+
+func TestPoolRemoveRemovesResourceFromPoolAndDoesNotStartNewCreationToMaintainMinSizeWhenPoolIsClosed(t *testing.T) {
+	createCallsChan := make(chan int, 4)
+	closeCallsChan := make(chan int, 4)
+
+	waitForRead := func(ch chan int) bool {
+		select {
+		case <-ch:
+			return true
+		case <-time.NewTimer(time.Second).C:
+			return false
+		}
+	}
+
+	var createCalls Counter
+	createFunc := func() (interface{}, error) {
+		n := createCalls.Next()
+		createCallsChan <- n
+		return n, nil
+	}
+
+	var closeCalls Counter
+	closeFunc := func(interface{}) error {
+		n := closeCalls.Next()
+		closeCallsChan <- n
+		return nil
+	}
+
+	pool := puddle.NewPool(createFunc, closeFunc)
+
+	// Ensure there are 2 resources available in pool
+	{
+		r1, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		r2, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		pool.Return(r1)
+		pool.Return(r2)
+	}
+
+	assert.Equal(t, 2, pool.Size())
+	pool.SetMinSize(2)
+	assert.Equal(t, 2, pool.Size())
+
+	{
+		r1, err := pool.Get(context.Background())
+		require.Nil(t, err)
+		r2, err := pool.Get(context.Background())
+		require.Nil(t, err)
+
+		pool.Close()
+
+		pool.Remove(r1)
+		pool.Remove(r2)
+	}
+
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(createCallsChan))
+	require.True(t, waitForRead(closeCallsChan))
+	require.True(t, waitForRead(closeCallsChan))
+
+	assert.Equal(t, 0, pool.Size())
+	assert.Equal(t, 2, createCalls.Value())
+	assert.Equal(t, 2, closeCalls.Value())
+}
+
 func TestPoolGetReturnsErrorWhenPoolIsClosed(t *testing.T) {
 	var createCalls Counter
 	createFunc := func() (interface{}, error) {