From 7df4afe7d3ae54415b939a88b42a8e5c5bcf3921 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Wed, 26 Dec 2018 13:01:31 -0600
Subject: [PATCH] Add Resource.Hijack

---
 pool.go      | 27 +++++++++++++++++++++++++++
 pool_test.go | 19 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/pool.go b/pool.go
index 4308e76..dd04126 100644
--- a/pool.go
+++ b/pool.go
@@ -10,6 +10,7 @@ const (
 	resourceStatusCreating  = 0
 	resourceStatusAvailable = iota
 	resourceStatusBorrowed  = iota
+	resourceStatusHijacked  = iota
 )
 
 const maxUint = ^uint(0)
@@ -32,13 +33,28 @@ func (res *Resource) Value() interface{} {
 }
 
 func (res *Resource) Release() {
+	if res.status != resourceStatusBorrowed {
+		panic("tried to release resource that is not acquired")
+	}
 	res.pool.releaseBorrowedResource(res)
 }
 
 func (res *Resource) Destroy() {
+	if res.status != resourceStatusBorrowed {
+		panic("tried to destroy resource that is not acquired")
+	}
 	res.pool.destroyBorrowedResource(res)
 }
 
+// Hijack removes the resource from the pool without destroying it. Caller is
+// responsible for cleanup of resource value.
+func (res *Resource) Hijack() {
+	if res.status != resourceStatusBorrowed {
+		panic("tried to hijack resource that is not acquired")
+	}
+	res.pool.hijackBorrowedResource(res)
+}
+
 // Pool is a thread-safe resource pool.
 type Pool struct {
 	cond       *sync.Cond
@@ -244,6 +260,17 @@ func (p *Pool) destroyBorrowedResource(res *Resource) {
 	p.cond.Signal()
 }
 
+func (p *Pool) hijackBorrowedResource(res *Resource) {
+	p.cond.L.Lock()
+
+	p.allResources = removeResource(p.allResources, res)
+	res.status = resourceStatusHijacked
+	p.destructWG.Done() // not responsible for destructing hijacked resources
+
+	p.cond.L.Unlock()
+	p.cond.Signal()
+}
+
 func removeResource(slice []*Resource, res *Resource) []*Resource {
 	for i := range slice {
 		if slice[i] == res {
diff --git a/pool_test.go b/pool_test.go
index e7b07c4..14711a4 100644
--- a/pool_test.go
+++ b/pool_test.go
@@ -264,6 +264,25 @@ func TestPoolCloseBlocksUntilAllResourcesReleasedAndClosed(t *testing.T) {
 }
 
 func TestResourceDestroyRemovesResourceFromPool(t *testing.T) {
+	createFunc, _ := createCreateResourceFunc()
+	var closeCalls Counter
+	closeFunc := func(interface{}) {
+		closeCalls.Next()
+	}
+
+	pool := puddle.NewPool(createFunc, closeFunc)
+
+	res, err := pool.Acquire(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, 1, res.Value())
+
+	res.Hijack()
+
+	assert.Equal(t, 0, pool.Size())
+	assert.Equal(t, 0, closeCalls.Value())
+}
+
+func TestResourceHijackRemovesResourceFromPoolButDoesNotDestroy(t *testing.T) {
 	createFunc, _ := createCreateResourceFunc()
 	pool := puddle.NewPool(createFunc, stubCloseRes)