From f63192c063d5612500c03782496ae01372faab50 Mon Sep 17 00:00:00 2001
From: James Hartig <me@jameshartig.com>
Date: Sat, 27 Aug 2022 11:30:19 -0400
Subject: [PATCH] Fix race with background Acquire creation

---
 pool.go      | 45 +++++++++++++++------------------------------
 pool_test.go |  3 +++
 2 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/pool.go b/pool.go
index a413e4d..fa2f86f 100644
--- a/pool.go
+++ b/pool.go
@@ -323,49 +323,34 @@ func (p *Pool[T]) Acquire(ctx context.Context) (*Resource[T], error) {
 					p.allResources = removeResource(p.allResources, res)
 					p.destructWG.Done()
 
-					select {
-					case <-ctx.Done():
-						if err == ctx.Err() {
-							p.canceledAcquireCount += 1
-						}
-					default:
-					}
-
-					p.cond.L.Unlock()
-					p.cond.Signal()
-
-					// try to notify the caller that we failed
+					// we can't use default here in case we get here before the caller is
+					// in the select
 					select {
 					case constructErrCh <- err:
-					default:
+					case <-ctx.Done():
+						p.canceledAcquireCount += 1
 					}
+					p.cond.L.Unlock()
+					p.cond.Signal()
 					return
 				}
 				res.value = value
 
-				// check the context now so we don't increment the metrics when the caller
-				// has already been cancelled
+				// assume that we will acquire it
+				res.status = resourceStatusAcquired
+				// we can't use default here in case we get here before the caller is
+				// in the select
 				select {
-				case <-ctx.Done():
-					p.cond.L.Unlock()
-				default:
-					// assume that we will acquire it
-					res.status = resourceStatusAcquired
-					// we have to increment these BEFORE the next select because otherwise
-					// they could run after Acquire has returned and that will mess up
-					// tests but this also means that these could be incremented even if
-					// the acquire times out, but we just checked that so the chances are
-					// slim
+				case constructErrCh <- nil:
 					p.emptyAcquireCount += 1
 					p.acquireCount += 1
 					p.acquireDuration += time.Duration(nanotime() - startNano)
 					p.cond.L.Unlock()
+					// we don't call Signal here we didn't change any of the resource pools
+				case <-ctx.Done():
+					p.canceledAcquireCount += 1
+					p.cond.L.Unlock()
 					// we don't call Signal here we didn't change any of the resopurce pools
-				}
-
-				select {
-				case constructErrCh <- nil:
-				default:
 					// since we couldn't send the constructed resource to the acquire
 					// function that means the caller has stopped waiting and we should
 					// just put this resource back in the pool
diff --git a/pool_test.go b/pool_test.go
index d7e82b7..e264f6e 100644
--- a/pool_test.go
+++ b/pool_test.go
@@ -648,6 +648,9 @@ func TestPoolStatCanceledAcquireDuringCreate(t *testing.T) {
 	_, err := pool.Acquire(ctx)
 	require.Equal(t, context.Canceled, err)
 
+	// sleep to give the constructor goroutine time to mark cancelled
+	time.Sleep(10 * time.Millisecond)
+
 	stat := pool.Stat()
 	assert.Equal(t, int64(0), stat.AcquireCount())
 	assert.Equal(t, int64(1), stat.CanceledAcquireCount())