From 7eab95c3a799f65d07c34af757173a5a0a44759a Mon Sep 17 00:00:00 2001 From: bestmountain Date: Tue, 30 Apr 2024 12:20:46 +0800 Subject: [PATCH 1/3] 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 From 66b44312a3ca6653016dcdcc21c9d6acb1c2f821 Mon Sep 17 00:00:00 2001 From: bestmountain Date: Tue, 30 Apr 2024 13:12:06 +0800 Subject: [PATCH 2/3] move MustGet to must.go --- must.go | 9 +++++++++ utils.go | 9 --------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/must.go b/must.go index 6ca49e45..13ebbfe3 100644 --- a/must.go +++ b/must.go @@ -1155,3 +1155,12 @@ 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 +} diff --git a/utils.go b/utils.go index e5e525ea..bacf5025 100644 --- a/utils.go +++ b/utils.go @@ -92,15 +92,6 @@ func NewPagePool(limit int) PagePool { return pp } -// 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 -} - // 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 From d60d5d53d66dbe2c4da082f55a06ede8f690be8b Mon Sep 17 00:00:00 2001 From: bestmountain Date: Tue, 30 Apr 2024 13:22:52 +0800 Subject: [PATCH 3/3] PagePool would not stuck while pp is not full --- page_test.go | 15 ++++++++++++++- utils.go | 11 +++++++---- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/page_test.go b/page_test.go index 1778213e..c88ad2de 100644 --- a/page_test.go +++ b/page_test.go @@ -1011,9 +1011,15 @@ func TestPagePool_Get_Negative(t *testing.T) { failContext, cancel := context.WithCancel(g.Context()) g.browser = g.browser.Context(failContext) // manipulate browser canceled by another thread - cancel() 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 { @@ -1022,6 +1028,13 @@ func TestPagePool_Get_Negative(t *testing.T) { 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 { diff --git a/utils.go b/utils.go index bacf5025..606e30a2 100644 --- a/utils.go +++ b/utils.go @@ -110,12 +110,15 @@ func (pp PagePool) Put(p *Page) { pp <- p } -// Cleanup helper, may stuck while PagePool is not full +// 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: } } }