diff --git a/allocate.go b/allocate.go index 2d2887d..ee94f14 100644 --- a/allocate.go +++ b/allocate.go @@ -157,6 +157,7 @@ func (p *ExecAllocator) Allocate(ctx context.Context) (*Browser, error) { if err != nil { return nil, err } + browser.process = cmd.Process browser.userDataDir = dataDir return browser, nil } diff --git a/browser.go b/browser.go index 837bd7d..3aa6b87 100644 --- a/browser.go +++ b/browser.go @@ -10,6 +10,7 @@ import ( "context" "encoding/json" "log" + "os" "sync/atomic" "github.com/mailru/easyjson" @@ -24,8 +25,6 @@ import ( // the browser process runner, WebSocket clients, associated targets, and // network, page, and DOM events. type Browser struct { - userDataDir string - conn Transport // next is the next message id. @@ -43,6 +42,16 @@ type Browser struct { // logging funcs logf func(string, ...interface{}) errf func(string, ...interface{}) + + // The optional fields below are helpful for some tests. + + // process can be initialized by the allocators which start a process + // when allocating a browser. + process *os.Process + + // userDataDir can be initialized by the allocators which set up user + // data dirs directly. + userDataDir string } type newTab struct { @@ -164,9 +173,7 @@ type tabEvent struct { func (b *Browser) run(ctx context.Context) { defer b.conn.Close() - // add cancel to context - ctx, cancel := context.WithCancel(ctx) - defer cancel() + cancel := FromContext(ctx).cancel // tabEventQueue is the queue of incoming target events, to be routed by // their session ID. @@ -179,10 +186,13 @@ func (b *Browser) run(ctx context.Context) { // connection. The separate goroutine is needed since a websocket read // is blocking, so it cannot be used in a select statement. go func() { - defer cancel() for { msg, err := b.conn.Read() if err != nil { + // If the websocket failed, most likely Chrome + // was closed or crashed. Cancel the entire + // Browser context to stop all activity. + cancel() return } if msg.Method == cdproto.EventRuntimeExceptionThrown { @@ -232,8 +242,6 @@ func (b *Browser) run(ctx context.Context) { // This goroutine handles tabs, as well as routing events to each tab // via the pages map. go func() { - defer cancel() - // This map is only safe for use within this goroutine, so don't // declare it as a Browser field. pages := make(map[target.SessionID]*Target, 1024) diff --git a/context.go b/context.go index 08d824a..1ccb8c6 100644 --- a/context.go +++ b/context.go @@ -30,6 +30,12 @@ type Context struct { // have its own unique Target pointing to a separate browser tab (page). Target *Target + // cancel simply cancels the context that was used to start Browser. + // This is useful to stop all activity and avoid deadlocks if we detect + // that the browser was closed or happened to crash. Note that this + // cancel function doesn't do any waiting. + cancel func() + // first records whether this context was the one that allocated // Browser. This is important, because its cancellation will stop the // entire browser handler, meaning that no further actions can be @@ -44,7 +50,7 @@ type Context struct { func NewContext(parent context.Context, opts ...ContextOption) (context.Context, context.CancelFunc) { ctx, cancel := context.WithCancel(parent) - c := &Context{} + c := &Context{cancel: cancel} if pc := FromContext(parent); pc != nil { c.Allocator = pc.Allocator c.Browser = pc.Browser diff --git a/context_test.go b/context_test.go index 402d097..fa6ae3a 100644 --- a/context_test.go +++ b/context_test.go @@ -2,7 +2,13 @@ package chromedp import ( "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "runtime" "testing" + "time" ) func TestTargets(t *testing.T) { @@ -48,3 +54,47 @@ func TestTargets(t *testing.T) { t.Fatal(err) } } + +func TestBrowserQuit(t *testing.T) { + t.Parallel() + if runtime.GOOS == "windows" { + t.Skip("os.Interrupt isn't supported on Windows") + } + + // Simulate a scenario where we navigate to a page that's slow to + // respond, and the browser is closed before we can finish the + // navigation. + serve := make(chan bool, 1) + close := make(chan bool, 1) + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close <- true + <-serve + fmt.Fprintf(w, "response") + })) + defer s.Close() + + ctx, cancel := NewContext(context.Background()) + defer cancel() + if err := Run(ctx); err != nil { + t.Fatal(err) + } + + go func() { + <-close + b := FromContext(ctx).Browser + if err := b.process.Signal(os.Interrupt); err != nil { + t.Error(err) + } + serve <- true + }() + + // Run should error with something other than "deadline exceeded" in + // much less than 5s. + ctx2, _ := context.WithTimeout(ctx, 5*time.Second) + switch err := Run(ctx2, Navigate(s.URL)); err { + case nil: + t.Fatal("did not expect a nil error") + case context.DeadlineExceeded: + t.Fatalf("did not expect a standard context error: %v", err) + } +}