Feat: Add the PagePool.TryGet method while it returns an error for pr…

…ocessing. (#1051) * PagePool add TryGet to deal error that Page() returns. * modify tryget * Add BrowserPool TryGet test * test coverage * go fmt * go fmt * pool.Get -> pool.MustGet; pool.Get allow error --------- Co-authored-by: bestmountain <zxcltj@foxmail.com>
go-rod · May 6, 2024 · 2d3e61f · 2d3e61f
1 parent 5a690b3
commit 2d3e61f
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 19 deletions.
diff --git a/browser_test.go b/browser_test.go
@@ -1,6 +1,7 @@
 package rod_test
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -442,13 +443,64 @@ func TestBrowserConnectFailure(t *testing.T) {
 func TestBrowserPool(_ *testing.T) {
 	pool := rod.NewBrowserPool(3)
 	create := func() *rod.Browser { return rod.New().MustConnect() }
-	b := pool.Get(create)
+	b := pool.MustGet(create)
 	pool.Put(b)
 	pool.Cleanup(func(p *rod.Browser) {
 		p.MustClose()
 	})
 }
 
+func TestBrowserPool_TryGet(t *testing.T) {
+	pool := rod.NewBrowserPool(3)
+	defer pool.Cleanup(func(p *rod.Browser) {
+		err := p.Close()
+		if err != nil {
+			t.Log(err)
+		}
+	})
+	create := func() (*rod.Browser, error) {
+		b := rod.New()
+		err := b.Connect()
+		return b, err
+	}
+	for i := 0; i < 4; i++ {
+		b, err := pool.Get(create)
+		if err != nil {
+			t.Fatal(err)
+		}
+		pool.Put(b)
+	}
+}
+
+func TestBrowserPool_TryGet_Negative(t *testing.T) {
+	pool := rod.NewBrowserPool(3)
+	defer pool.Cleanup(func(p *rod.Browser) {
+		err := p.Close()
+		if err != nil {
+			t.Log(err)
+		}
+	})
+	ctx, cancel := context.WithCancel(context.Background())
+	create := func() (*rod.Browser, error) {
+		b := rod.New().Context(ctx)
+		err := b.Connect()
+		return b, err
+	}
+	b, err := pool.Get(create)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pool.Put(b)
+	cancel()
+	b, err = pool.Get(create)
+	if err != nil {
+		t.Log(err)
+	} else {
+		pool.Put(b)
+		t.FailNow()
+	}
+}
+
 func TestOldBrowser(t *testing.T) {
 	t.Skip()
 

diff --git 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.
@@ -632,7 +632,7 @@ func ExampleBrowser_pool() {
 			defer wg.Done()
 
 			// Get a browser instance from the pool
-			browser := pool.Get(create)
+			browser := 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
@@ -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/must.go b/must.go
@@ -1155,3 +1155,21 @@ func (el *Element) MustGetXPath(optimized bool) string {
 	el.e(err)
 	return xpath
 }
+
+// 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()
+	}
+	return p
+}
+
+// MustGet a browser from the pool. Use the [BrowserPool.Put] to make it reusable later.
+func (bp BrowserPool) MustGet(create func() *Browser) *Browser {
+	p := <-bp
+	if p == nil {
+		p = create()
+	}
+	return p
+}
diff --git a/page_test.go b/page_test.go
@@ -978,13 +978,73 @@ func TestPagePool(t *testing.T) {
 
 	pool := rod.NewPagePool(3)
 	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) {
 		p.MustClose()
 	})
 }
 
+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: ""})
+	}
+	for i := 0; i < 4; i++ {
+		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
+	pool := rod.NewPagePool(3)
+
+	defer pool.Cleanup(func(p *rod.Page) {
+		err := p.Close()
+		if err != nil {
+			t.Log(err)
+		}
+	})
+
+	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)
+
+	cancel()
+	p, err = pool.Get(create)
+	if err != nil {
+		t.Log(err)
+	} else {
+		pool.Put(p)
+		t.FailNow()
+	}
+}
+
 func TestPageUseNonExistSession(t *testing.T) {
 	g := setup(t)
 

diff --git a/utils.go b/utils.go
@@ -92,13 +92,13 @@ 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 {
+// 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
 	if p == nil {
-		p = create()
+		return create()
 	}
-	return p
+	return p, nil
 }
 
 // Put a page back to the pool.
@@ -109,9 +109,12 @@ func (pp PagePool) Put(p *Page) {
 // Cleanup helper.
 func (pp PagePool) Cleanup(iteratee func(*Page)) {
 	for i := 0; i < cap(pp); i++ {
-		p := <-pp
-		if p != nil {
-			iteratee(p)
+		select {
+		case p := <-pp:
+			if p != nil {
+				iteratee(p)
+			}
+		default:
 		}
 	}
 }
@@ -131,13 +134,13 @@ func NewBrowserPool(limit int) BrowserPool {
 	return pp
 }
 
-// Get a browser from the pool. Use the [BrowserPool.Put] to make it reusable later.
-func (bp BrowserPool) Get(create func() *Browser) *Browser {
+// Get a browser from the pool, allow error. Use the [BrowserPool.Put] to make it reusable later.
+func (bp BrowserPool) Get(create func() (*Browser, error)) (*Browser, error) {
 	p := <-bp
 	if p == nil {
-		p = create()
+		return create()
 	}
-	return p
+	return p, nil
 }
 
 // Put a browser back to the pool.
@@ -148,9 +151,12 @@ func (bp BrowserPool) Put(p *Browser) {
 // Cleanup helper.
 func (bp BrowserPool) Cleanup(iteratee func(*Browser)) {
 	for i := 0; i < cap(bp); i++ {
-		p := <-bp
-		if p != nil {
-			iteratee(p)
+		select {
+		case p := <-bp:
+			if p != nil {
+				iteratee(p)
+			}
+		default:
 		}
 	}
 }