make sure closed page doesn't block any action

browser.go

@@ -250,12 +250,14 @@ func (b *Browser) Call(ctx context.Context, sessionID, methodName string, params
 
 // PageFromSession is used for low-level debugging
 func (b *Browser) PageFromSession(sessionID proto.TargetSessionID) *Page {
+	sessionCtx, cancel := context.WithCancel(b.ctx)
 	return &Page{
-		e:         b.e,
-		ctx:       b.ctx,
-		sleeper:   b.sleeper,
-		browser:   b,
-		SessionID: sessionID,
+		e:             b.e,
+		ctx:           sessionCtx,
+		sessionCancel: cancel,
+		sleeper:       b.sleeper,
+		browser:       b,
+		SessionID:     sessionID,
 	}
 }
 
@@ -277,17 +279,20 @@ func (b *Browser) PageFromTarget(targetID proto.TargetTargetID) (*Page, error) {
 		return nil, err
 	}
 
+	sessionCtx, cancel := context.WithCancel(b.ctx)
+
 	page = &Page{
-		e:           b.e,
-		ctx:         b.ctx,
-		sleeper:     b.sleeper,
-		browser:     b,
-		TargetID:    targetID,
-		SessionID:   session.SessionID,
-		FrameID:     proto.PageFrameID(targetID),
-		jsCtxLock:   &sync.Mutex{},
-		jsCtxID:     new(proto.RuntimeRemoteObjectID),
-		helpersLock: &sync.Mutex{},
+		e:             b.e,
+		ctx:           sessionCtx,
+		sessionCancel: cancel,
+		sleeper:       b.sleeper,
+		browser:       b,
+		TargetID:      targetID,
+		SessionID:     session.SessionID,
+		FrameID:       proto.PageFrameID(targetID),
+		jsCtxLock:     &sync.Mutex{},
+		jsCtxID:       new(proto.RuntimeRemoteObjectID),
+		helpersLock:   &sync.Mutex{},
 	}
 
 	page.root = page
@@ -304,6 +309,8 @@ func (b *Browser) PageFromTarget(targetID proto.TargetTargetID) (*Page, error) {
 
 	b.cachePage(page)
 
+	page.initEvents()
+
 	// If we don't enable it, it will cause a lot of unexpected browser behavior.
 	// Such as proto.PageAddScriptToEvaluateOnNewDocument won't work.
 	page.EnableDomain(&proto.PageEnable{})
@@ -408,10 +415,9 @@ func (b *Browser) eachEvent(sessionID proto.TargetSessionID, callbacks ...interf
 
 // Event of the browser
 func (b *Browser) Event() <-chan *Message {
-	src := b.event.Subscribe()
+	src := b.event.Subscribe(b.ctx)
 	dst := make(chan *Message)
 	go func() {
-		defer b.event.Unsubscribe(src)
 		defer close(dst)
 		for {
 			select {
@@ -433,10 +439,9 @@ func (b *Browser) Event() <-chan *Message {
 }
 
 func (b *Browser) initEvents() {
-	b.event = goob.New()
+	b.event = goob.New(b.ctx)
 
 	go func() {
-		defer b.event.Close()
 		for e := range b.client.Event() {
 			b.event.Publish(&Message{
 				SessionID: proto.TargetSessionID(e.SessionID),

browser_test.go

@@ -1,7 +1,6 @@
 package rod_test
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -89,37 +88,35 @@ func TestPageFromTarget(t *testing.T) {
 func TestBrowserPages(t *testing.T) {
 	g := setup(t)
 
-	g.newPage(g.blank()).MustWaitLoad()
-
-	pages := g.browser.MustPages()
-
-	g.Len(pages, 3)
+	u := launcher.New().MustLaunch()
+	mc := newMockClient(u)
+	b := rod.New().Client(mc).MustConnect()
+	g.Cleanup(func() { b.MustClose() })
+	b.MustPage().MustWaitLoad()
+	pages := b.MustPages()
+	g.Gte(len(pages), 1)
 
 	{
-		g.mc.stub(1, proto.TargetGetTargets{}, func(send StubSend) (gson.JSON, error) {
+		mc.stub(1, proto.TargetGetTargets{}, func(send StubSend) (gson.JSON, error) {
 			d, _ := send()
 			return *d.Set("targetInfos.0.type", "iframe"), nil
 		})
-		pages := g.browser.MustPages()
-		g.Len(pages, 2)
+		b.MustPages()
 	}
 
 	g.Panic(func() {
-		g.mc.stubErr(1, proto.TargetCreateTarget{})
-		g.browser.MustPage()
+		mc.stubErr(1, proto.TargetCreateTarget{})
+		b.MustPage()
 	})
 	g.Panic(func() {
-		g.mc.stubErr(1, proto.TargetGetTargets{})
-		g.browser.MustPages()
+		mc.stubErr(1, proto.TargetGetTargets{})
+		b.MustPages()
 	})
 	g.Panic(func() {
-		res, err := proto.TargetCreateTarget{URL: "about:blank"}.Call(g.browser)
+		_, err := proto.TargetCreateTarget{URL: "about:blank"}.Call(b)
 		g.E(err)
-		defer func() {
-			g.browser.MustPageFromTargetID(res.TargetID).MustClose()
-		}()
-		g.mc.stubErr(1, proto.TargetAttachToTarget{})
-		g.browser.MustPages()
+		mc.stubErr(1, proto.TargetAttachToTarget{})
+		b.MustPages()
 	})
 }
 
@@ -146,21 +143,6 @@ func TestBrowserEvent(t *testing.T) {
 	<-wait
 }
 
-func TestBrowserEventClose(t *testing.T) {
-	event := make(chan *cdp.Event)
-	c := &MockClient{
-		connect: func() error { return nil },
-		call: func(ctx context.Context, sessionID, method string, params interface{}) ([]byte, error) {
-			return nil, errors.New("err")
-		},
-		event: event,
-	}
-	b := rod.New().Client(c)
-	_ = b.Connect()
-	b.Event()
-	close(event)
-}
-
 func TestBrowserWaitEvent(t *testing.T) {
 	g := setup(t)
 

go.mod

@@ -3,8 +3,8 @@ module github.com/go-rod/rod
 go 1.16
 
 require (
-	github.com/ysmood/goob v0.3.1
-	github.com/ysmood/got v0.23.0
+	github.com/ysmood/goob v0.4.0
+	github.com/ysmood/got v0.23.2
 	github.com/ysmood/gotrace v0.6.0
 	github.com/ysmood/gson v0.7.1
 	github.com/ysmood/leakless v0.7.0

go.sum

@@ -1,7 +1,7 @@
-github.com/ysmood/goob v0.3.1 h1:qMp5364BGS1DLJVrAqUxTF6KOFt0YDot8GC70u/0jbI=
-github.com/ysmood/goob v0.3.1/go.mod h1:S3lq113Y91y1UBf1wj1pFOxeahvfKkCk6mTWTWbDdWs=
-github.com/ysmood/got v0.23.0 h1:WUk6abT2QPjTsi/f0NpIlkpXUVla4tsPPDBjytFDWBw=
-github.com/ysmood/got v0.23.0/go.mod h1:pE1l4LOwOBhQg6A/8IAatkGp7uZjnalzrZolnlhhMgY=
+github.com/ysmood/goob v0.4.0 h1:HsxXhyLBeGzWXnqVKtmT9qM7EuVs/XOgkX7T6r1o1AQ=
+github.com/ysmood/goob v0.4.0/go.mod h1:u6yx7ZhS4Exf2MwciFr6nIM8knHQIE22lFpWHnfql18=
+github.com/ysmood/got v0.23.2 h1:U2U0vyQ/gDaawkKJZK/hyza8UUXbWCurbmazK7AcAfY=
+github.com/ysmood/got v0.23.2/go.mod h1:pE1l4LOwOBhQg6A/8IAatkGp7uZjnalzrZolnlhhMgY=
 github.com/ysmood/gotrace v0.6.0 h1:SyI1d4jclswLhg7SWTL6os3L1WOKeNn/ZtzVQF8QmdY=
 github.com/ysmood/gotrace v0.6.0/go.mod h1:TzhIG7nHDry5//eYZDYcTzuJLYQIkykJzCRIo4/dzQM=
 github.com/ysmood/gson v0.7.1 h1:zKL2MTGtynxdBdlZjyGsvEOZ7dkxaY5TH6QhAbTgz0Q=

page.go

@@ -13,6 +13,7 @@ import (
 	"github.com/go-rod/rod/lib/js"
 	"github.com/go-rod/rod/lib/proto"
 	"github.com/go-rod/rod/lib/utils"
+	"github.com/ysmood/goob"
 	"github.com/ysmood/gson"
 )
 
@@ -21,8 +22,9 @@ var _ proto.Client = &Page{}
 var _ proto.Contextable = &Page{}
 var _ proto.Sessionable = &Page{}
 
-// Page represents the webpage
-// We try to hold as less states as possible
+// Page represents the webpage.
+// We try to hold as less states as possible.
+// When a page is closed by Rod or not all the ongoing operations an events on it will abort.
 type Page struct {
 	// TargetID is a unique ID for a remote page.
 	// It's usually used in events sent from the browser to tell which page an event belongs to.
@@ -42,11 +44,15 @@ type Page struct {
 
 	ctx context.Context
 
+	// Used to abort all ongoing actions when a page closes.
+	sessionCancel func()
+
 	root *Page
 
 	sleeper func() utils.Sleeper
 
 	browser *Browser
+	event   *goob.Observable
 
 	// devices
 	Mouse    *Mouse
@@ -282,7 +288,7 @@ func (p *Page) StopLoading() error {
 	return proto.PageStopLoading{}.Call(p)
 }
 
-// Close tries to close page, running its beforeunload hooks, if any.
+// Close tries to close page, running its beforeunload hooks, if has any.
 func (p *Page) Close() error {
 	p.browser.targetsLock.Lock()
 	defer p.browser.targetsLock.Unlock()
@@ -708,27 +714,49 @@ func (p *Page) Call(ctx context.Context, sessionID, methodName string, params in
 // Event of the page
 func (p *Page) Event() <-chan *Message {
 	dst := make(chan *Message)
-	p, cancel := p.WithCancel()
-	messages := p.browser.Context(p.ctx).Event()
 
 	go func() {
 		defer close(dst)
-		defer cancel()
-		for m := range messages {
-			detached := proto.TargetDetachedFromTarget{}
-			if m.Load(&detached) && detached.SessionID == p.SessionID {
+		s := p.event.Subscribe(p.ctx)
+		for {
+			select {
+			case <-p.ctx.Done():
 				return
-			}
-
-			if m.SessionID == p.SessionID {
+			case msg, ok := <-s:
+				if !ok {
+					return
+				}
 				select {
 				case <-p.ctx.Done():
 					return
-				case dst <- m:
+				case dst <- msg.(*Message):
 				}
 			}
 		}
 	}()
 
 	return dst
 }
+
+func (p *Page) initEvents() {
+	p.event = goob.New(p.ctx)
+
+	go func() {
+		for msg := range p.browser.Context(p.ctx).Event() {
+			detached := proto.TargetDetachedFromTarget{}
+			destroyed := proto.TargetTargetDestroyed{}
+
+			if (msg.Load(&detached) && detached.SessionID == p.SessionID) ||
+				(msg.Load(destroyed) && destroyed.TargetID == p.TargetID) {
+				p.sessionCancel()
+				return
+			}
+
+			if msg.SessionID != p.SessionID {
+				continue
+			}
+
+			p.event.Publish(msg)
+		}
+	}()
+}

page_test.go

@@ -943,3 +943,14 @@ func TestPageElementFromObjectErr(t *testing.T) {
 	g.mc.stubErr(1, proto.RuntimeEvaluate{})
 	g.Err(p.ElementFromObject(obj.Object))
 }
+
+func TestPageActionAfterClose(t *testing.T) {
+	g := setup(t)
+
+	p := g.browser.MustPage(g.blank())
+
+	p.MustClose()
+
+	_, err := p.Element("nonexists")
+	g.Is(err, context.Canceled)
+}

query_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/go-rod/rod"
 	"github.com/go-rod/rod/lib/cdp"
 	"github.com/go-rod/rod/lib/defaults"
+	"github.com/go-rod/rod/lib/launcher"
 	"github.com/go-rod/rod/lib/proto"
 	"github.com/go-rod/rod/lib/utils"
 	"github.com/ysmood/gson"
@@ -24,11 +25,16 @@ func TestPageElements(t *testing.T) {
 	g.Eq("submit", list.Last().MustText())
 }
 
-func TestPages(t *testing.T) {
+func TestPagesQuery(t *testing.T) {
 	g := setup(t)
 
-	g.page.MustNavigate(g.srcFile("fixtures/click.html")).MustWaitLoad()
-	pages := g.browser.MustPages()
+	u := launcher.New().MustLaunch()
+	mc := newMockClient(u)
+	b := rod.New().Client(mc).MustConnect()
+	g.Cleanup(func() { b.MustClose() })
+
+	b.MustPage(g.srcFile("fixtures/click.html")).MustWaitLoad()
+	pages := b.MustPages()
 
 	g.True(pages.MustFind("button").MustHas("button"))
 	g.Panic(func() { rod.Pages{}.MustFind("____") })
@@ -43,15 +49,28 @@ func TestPages(t *testing.T) {
 	})
 
 	g.Panic(func() {
-		g.mc.stubErr(1, proto.RuntimeCallFunctionOn{})
+		mc.stubErr(1, proto.RuntimeCallFunctionOn{})
 		pages.MustFind("button")
 	})
 	g.Panic(func() {
-		g.mc.stubErr(1, proto.RuntimeCallFunctionOn{})
+		mc.stubErr(1, proto.RuntimeCallFunctionOn{})
 		pages.MustFindByURL("____")
 	})
 }
 
+func TestPagesOthers(t *testing.T) {
+	g := setup(t)
+
+	list := rod.Pages{}
+	g.Nil(list.First())
+	g.Nil(list.Last())
+
+	list = append(list, &rod.Page{})
+
+	g.NotNil(list.First())
+	g.NotNil(list.Last())
+}
+
 func TestPageHas(t *testing.T) {
 	g := setup(t)
 
@@ -401,16 +420,3 @@ func TestElementsOthers(t *testing.T) {
 	g.Nil(list.First())
 	g.Nil(list.Last())
 }
-
-func TestPagesOthers(t *testing.T) {
-	g := setup(t)
-
-	list := rod.Pages{}
-	g.Nil(list.First())
-	g.Nil(list.Last())
-
-	list = append(list, &rod.Page{})
-
-	g.NotNil(list.First())
-	g.NotNil(list.Last())
-}