From 7eab95c3a799f65d07c34af757173a5a0a44759a Mon Sep 17 00:00:00 2001
From: bestmountain <zxcltj@foxmail.com>
Date: Tue, 30 Apr 2024 12:20:46 +0800
Subject: [PATCH] replace PagePool Get to MustGet, add Get method allow error
 return.

---
 examples_test.go                              |  2 +-
 .../use-rod-like-chrome-extension/main.go     |  2 +-
 page_test.go                                  | 49 ++++++++++++++++++-
 utils.go                                      | 19 +++++--
 4 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/examples_test.go b/examples_test.go
index 3335a489..471a555d 100644
--- a/examples_test.go
+++ b/examples_test.go
@@ -582,7 +582,7 @@ func ExamplePage_pool() {
 	}
 
 	yourJob := func() {
-		page := pool.Get(create)
+		page := pool.MustGet(create)
 
 		// Put the instance back to the pool after we're done,
 		// so the instance can be reused by other goroutines.
diff --git a/lib/examples/use-rod-like-chrome-extension/main.go b/lib/examples/use-rod-like-chrome-extension/main.go
index 158c42a1..9980bfc4 100644
--- a/lib/examples/use-rod-like-chrome-extension/main.go
+++ b/lib/examples/use-rod-like-chrome-extension/main.go
@@ -51,7 +51,7 @@ func linkPreviewer(browser *rod.Browser) {
 
 		// Expose a function to the page to provide preview
 		page.MustExpose("getPreview", func(url gson.JSON) (interface{}, error) {
-			p := pool.Get(create)
+			p := pool.MustGet(create)
 			defer pool.Put(p)
 			p.MustNavigate(url.Str())
 			return base64.StdEncoding.EncodeToString(p.MustScreenshot()), nil
diff --git a/page_test.go b/page_test.go
index 37e73757..1778213e 100644
--- a/page_test.go
+++ b/page_test.go
@@ -977,12 +977,57 @@ func TestPagePool(t *testing.T) {
 	g := setup(t)
 
 	pool := rod.NewPagePool(3)
+	defer pool.Cleanup(func(p *rod.Page) {
+		p.MustClose()
+	})
 	create := func() *rod.Page { return g.browser.MustPage() }
-	p := pool.Get(create)
+	p := pool.MustGet(create)
 	pool.Put(p)
-	pool.Cleanup(func(p *rod.Page) {
+}
+
+func TestPagePool_Get(t *testing.T) {
+	g := setup(t)
+
+	pool := rod.NewPagePool(3)
+	defer pool.Cleanup(func(p *rod.Page) {
 		p.MustClose()
 	})
+	create := func() (*rod.Page, error) {
+		b, err := g.browser.Incognito()
+		if err != nil {
+			return nil, err
+		}
+		return b.Page(proto.TargetCreateTarget{URL: ""})
+	}
+	p, err := pool.Get(create)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pool.Put(p)
+}
+
+func TestPagePool_Get_Negative(t *testing.T) {
+	g := setup(t)
+	failContext, cancel := context.WithCancel(g.Context())
+	g.browser = g.browser.Context(failContext)
+	// manipulate browser canceled by another thread
+	cancel()
+	pool := rod.NewPagePool(3)
+
+	create := func() (*rod.Page, error) {
+		b, err := g.browser.Incognito()
+		if err != nil {
+			return nil, err
+		}
+		return b.Page(proto.TargetCreateTarget{URL: ""})
+	}
+	p, err := pool.Get(create)
+	if err != nil {
+		t.Log(err)
+	} else {
+		pool.Put(p)
+		t.FailNow()
+	}
 }
 
 func TestPageUseNonExistSession(t *testing.T) {
diff --git a/utils.go b/utils.go
index a3b31740..e5e525ea 100644
--- a/utils.go
+++ b/utils.go
@@ -92,8 +92,8 @@ func NewPagePool(limit int) PagePool {
 	return pp
 }
 
-// Get a page from the pool. Use the [PagePool.Put] to make it reusable later.
-func (pp PagePool) Get(create func() *Page) *Page {
+// MustGet a page from the pool. Use the [PagePool.Put] to make it reusable later.
+func (pp PagePool) MustGet(create func() *Page) *Page {
 	p := <-pp
 	if p == nil {
 		p = create()
@@ -101,12 +101,25 @@ func (pp PagePool) Get(create func() *Page) *Page {
 	return p
 }
 
+// Get a page from the pool, allow error. Use the [PagePool.Put] to make it reusable later.
+func (pp PagePool) Get(create func() (*Page, error)) (*Page, error) {
+	p := <-pp
+	var err error
+	if p == nil {
+		p, err = create()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return p, nil
+}
+
 // Put a page back to the pool.
 func (pp PagePool) Put(p *Page) {
 	pp <- p
 }
 
-// Cleanup helper.
+// Cleanup helper, may stuck while PagePool is not full
 func (pp PagePool) Cleanup(iteratee func(*Page)) {
 	for i := 0; i < cap(pp); i++ {
 		p := <-pp