avoid hanging when Chrome is closed separately

It's Run that actually starts a Browser, not NewContext. If the browser
is closed or crashes, the browser handler will fail to read from the
websocket, and its goroutines will stop.

However, the target handler's goroutines may not stop. The browser
handler uses a separate cancel function to stop itself when encountering
a websocket error, so that doesn't propagate to the original context
children, like the target handler.

To fix this, make it so that NewContext can keep the cancel function
around, for Run to use it in this scenario. And add a test case that
tests this very edge case, which used to time out before the fix.

Fixes #289.

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
 }

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)

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

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)
+	}
+}