use generics for page pool and browser pool

go-rod · Jul 6, 2024 · be0491c · be0491c
1 parent 9142e29
commit be0491c
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 223 deletions.
diff --git a/browser_test.go b/browser_test.go
@@ -1,7 +1,6 @@
 package rod_test
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -17,6 +16,7 @@ import (
 	"github.com/go-rod/rod/lib/launcher"
 	"github.com/go-rod/rod/lib/proto"
 	"github.com/go-rod/rod/lib/utils"
+	"github.com/ysmood/got"
 	"github.com/ysmood/gson"
 )
 
@@ -440,65 +440,24 @@ func TestBrowserConnectFailure(t *testing.T) {
 	}
 }
 
-func TestBrowserPool(_ *testing.T) {
-	pool := rod.NewBrowserPool(3)
-	create := func() *rod.Browser { return rod.New().MustConnect() }
-	b := pool.MustGet(create)
-	pool.Put(b)
-	pool.Cleanup(func(p *rod.Browser) {
-		p.MustClose()
-	})
-}
+func TestBrowserPool(t *testing.T) {
+	g := got.T(t)
 
-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)
-		}
+	b, err := pool.Get(func() (*rod.Browser, error) {
+		browser := rod.New()
+		return browser, browser.Connect()
 	})
-	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)
-	}
+	g.E(err)
 	pool.Put(b)
-	cancel()
-	b, err = pool.Get(create)
-	if err != nil {
-		t.Log(err)
-	} else {
-		pool.Put(b)
-		t.FailNow()
-	}
+
+	b = pool.MustGet(func() *rod.Browser { return rod.New().MustConnect() })
+	pool.Put(b)
+
+	pool.Cleanup(func(p *rod.Browser) {
+		p.MustClose()
+	})
 }
 
 func TestOldBrowser(t *testing.T) {

diff --git a/must.go b/must.go
@@ -1156,20 +1156,11 @@ func (el *Element) MustGetXPath(optimized bool) string {
 	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()
+// MustGet an elem from the pool. Use the [Pool[T].Put] to make it reusable later.
+func (p Pool[T]) MustGet(create func() *T) *T {
+	elem := <-p
+	if elem == nil {
+		elem = 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
+	return elem
 }
diff --git a/page_test.go b/page_test.go
@@ -975,72 +975,19 @@ func TestPagePool(t *testing.T) {
 	g := setup(t)
 
 	pool := rod.NewPagePool(3)
-	create := func() *rod.Page { return g.browser.MustPage() }
-	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)
-		}
+	p, err := pool.Get(func() (*rod.Page, error) {
+		return g.browser.Page(proto.TargetCreateTarget{})
 	})
+	g.E(err)
+	pool.Put(p)
 
-	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)
-	}
+	p = pool.MustGet(func() *rod.Page { return g.browser.MustPage() })
 	pool.Put(p)
 
-	cancel()
-	p, err = pool.Get(create)
-	if err != nil {
-		t.Log(err)
-	} else {
-		pool.Put(p)
-		t.FailNow()
-	}
+	pool.Cleanup(func(p *rod.Page) {
+		p.MustClose()
+	})
 }
 
 func TestPageUseNonExistSession(t *testing.T) {

diff --git a/setup_test.go b/setup_test.go
@@ -38,7 +38,7 @@ func init() {
 	launcher.NewBrowser().MustGet() // preload browser to local
 }
 
-var testerPool TesterPool
+var testerPool rod.Pool[G]
 
 func TestMain(m *testing.M) {
 	testerPool = newTesterPool()
@@ -48,7 +48,9 @@ func TestMain(m *testing.M) {
 		os.Exit(code)
 	}
 
-	testerPool.cleanup()
+	testerPool.Cleanup(func(g *G) {
+		g.browser.MustClose()
+	})
 
 	if err := gotrace.Check(0, gotrace.IgnoreFuncs("internal/poll.runtime_pollWait")); err != nil {
 		log.Fatal(err)
@@ -74,34 +76,19 @@ type G struct {
 	cancelTimeout func()
 }
 
-// TesterPool if we don't use pool to cache, the total time will be much longer.
-type TesterPool struct {
-	pool     chan *G
-	parallel int
-}
-
-func newTesterPool() TesterPool {
+// If we don't use pool to cache, the total time will be much longer.
+func newTesterPool() rod.Pool[G] {
 	parallel := got.Parallel()
 	if parallel == 0 {
 		parallel = runtime.GOMAXPROCS(0)
 	}
 
 	fmt.Println("parallel test", parallel) //nolint: forbidigo
 
-	cp := TesterPool{
-		pool:     make(chan *G, parallel),
-		parallel: parallel,
-	}
-
-	for i := 0; i < parallel; i++ {
-		cp.pool <- nil
-	}
-
-	return cp
+	return rod.NewPool[G](parallel)
 }
 
-// new tester.
-func (tp TesterPool) new() *G {
+func newTester() *G {
 	u := launcher.New().Set("proxy-bypass-list", "<-loopback>").MustLaunch()
 
 	mc := newMockClient(u)
@@ -131,11 +118,8 @@ func setup(t *testing.T) G {
 		t.Parallel()
 	}
 
-	tester := <-testerPool.pool
-	if tester == nil {
-		tester = testerPool.new()
-	}
-	t.Cleanup(func() { testerPool.pool <- tester })
+	tester := testerPool.MustGet(newTester)
+	t.Cleanup(func() { testerPool.Put(tester) })
 
 	tester.G = got.New(t)
 	tester.mc.t = t
@@ -146,14 +130,6 @@ func setup(t *testing.T) G {
 	return *tester
 }
 
-func (tp TesterPool) cleanup() {
-	for i := 0; i < tp.parallel; i++ {
-		if t := <-testerPool.pool; t != nil {
-			t.browser.MustClose()
-		}
-	}
-}
-
 func (g G) enableCDPLog() {
 	g.mc.principal.Logger(rod.DefaultLogger)
 }