Close pool in Cleanup to prevent Get getting stuck in concurrent scenarios.

browser_test.go

@@ -460,6 +460,27 @@ func TestBrowserPool(t *testing.T) {
 	})
 }
 
+func TestBrowserPoolNotStuck(t *testing.T) {
+	g := got.T(t)
+
+	pool := rod.NewBrowserPool(3)
+
+	pool.Cleanup(func(p *rod.Browser) {
+		p.MustClose()
+	})
+
+	b, err := pool.Get(func() (*rod.Browser, error) {
+		browser := rod.New()
+		return browser, browser.Connect()
+	})
+	if err != nil {
+		g.Equal(err, errors.New("pool has been cleaned up"))
+	} else {
+		pool.Put(b) // will panic
+		g.Fail()
+	}
+}
+
 func TestOldBrowser(t *testing.T) {
 	t.Skip()
 

page_test.go

@@ -992,6 +992,25 @@ func TestPagePool(t *testing.T) {
 	})
 }
 
+func TestPagePoolNotStuck(t *testing.T) {
+	g := setup(t)
+	pool := rod.NewPagePool(3)
+
+	pool.Cleanup(func(p *rod.Page) {
+		p.MustClose()
+	})
+
+	b, err := pool.Get(func() (*rod.Page, error) {
+		return g.browser.Page(proto.TargetCreateTarget{})
+	})
+	if err != nil {
+		g.Equal(err, errors.New("pool has been cleaned up"))
+	} else {
+		pool.Put(b) // will panic
+		g.Fail()
+	}
+}
+
 func TestPageUseNonExistSession(t *testing.T) {
 	g := setup(t)
 

utils.go

@@ -5,6 +5,7 @@
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -102,16 +103,20 @@
 	return p
 }
 
-// Get a elem from the pool, allow error. Use the [Pool[T].Put] to make it reusable later.
+// Get an elem from the pool, allow error. Use the [Pool[T].Put] to make it reusable later.
 func (p Pool[T]) Get(create func() (*T, error)) (elem *T, err error) {
-	elem = <-p
-	if elem == nil {
-		elem, err = create()
+	// use range to prevent getting stuck in concurrent scenarios
+	for elem = range p {
+		if elem == nil {
+			elem, err = create()
+		}
+		return
 	}
-	return
+	// here p has been closed
+	return nil, errors.New("pool has been cleaned up")
 }
 
-// Put an elem back to the pool.
+// Put an elem back to the pool. Should not be called when Get returns Error
 func (p Pool[T]) Put(elem *T) {
 	p <- elem
 }
@@ -127,6 +132,8 @@
 		default:
 		}
 	}
+	// close Pool after iterating to prevent Get getting stuck
+	close(p)
 }
 
 var _ io.ReadCloser = &StreamReader{}