From 487e8f2bf221253f703cccce47c719cd34711bf2 Mon Sep 17 00:00:00 2001
From: bestmountain <zxcltj@foxmail.com>
Date: Sun, 29 Sep 2024 11:44:01 +0800
Subject: [PATCH] close pool in Cleanup to prevent Get getting stuck in
 concurrent scenarios

---
 browser_test.go | 21 +++++++++++++++++++++
 page_test.go    | 19 +++++++++++++++++++
 utils.go        | 19 +++++++++++++------
 3 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/browser_test.go b/browser_test.go
index 419ebd3e..f0cc1dd4 100644
--- a/browser_test.go
+++ b/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()
 
diff --git a/page_test.go b/page_test.go
index 0bb3fdda..9f1f9e3d 100644
--- a/page_test.go
+++ b/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)
 
diff --git a/utils.go b/utils.go
index cae40b11..1e938952 100644
--- a/utils.go
+++ b/utils.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -102,16 +103,20 @@ func NewPool[T any](limit int) Pool[T] {
 	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 @@ func (p Pool[T]) Cleanup(iteratee func(*T)) {
 		default:
 		}
 	}
+	// close Pool after iterating to prevent Get getting stuck
+	close(p)
 }
 
 var _ io.ReadCloser = &StreamReader{}