diff --git a/browser_test.go b/browser_test.go index 38aecfd4..37083729 100644 --- 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 index 3335a489..db8d0475 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. @@ -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 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/must.go b/must.go index 6ca49e45..dea2c6a3 100644 --- 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 index 37e73757..810e1706 100644 --- 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 index a3b31740..8fa76d1b 100644 --- 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: } } }