wait for cleanup when cancelling the first context

We were doing this for extra open tabs, since NewContext was taking care
of detaching and closing the respective target.

And cancelling an entire allocator also properly waited for all the
resources, such as processes and temporary directories, to be cleaned
up.

However, this didn't work for a browser's first context; cancelling it
should wait for that one browser's resources to be cleaned up, but that
wasn't implemented. Do that, fixing the TODO in TestExecAllocator.

allocate.go

@@ -82,6 +82,11 @@ type ExecAllocator struct {
 
 // Allocate satisfies the Allocator interface.
 func (p *ExecAllocator) Allocate(ctx context.Context) (*Browser, error) {
+	c := FromContext(ctx)
+	if c == nil {
+		return nil, ErrInvalidContext
+	}
+
 	var args []string
 	for name, value := range p.initFlags {
 		switch value := value.(type) {
@@ -110,7 +115,8 @@ func (p *ExecAllocator) Allocate(ctx context.Context) (*Browser, error) {
 	args = append(args, "--remote-debugging-port=0")
 
 	var cmd *exec.Cmd
-	p.wg.Add(1)
+	p.wg.Add(1) // for the entire allocator
+	c.wg.Add(1) // for this browser's root context
 	go func() {
 		<-ctx.Done()
 		// First wait for the process to be finished.
@@ -122,6 +128,7 @@ func (p *ExecAllocator) Allocate(ctx context.Context) (*Browser, error) {
 			os.RemoveAll(dataDir)
 		}
 		p.wg.Done()
+		c.wg.Done()
 	}()
 
 	// force the first page to be blank, instead of the welcome page

allocate_test.go

@@ -30,8 +30,7 @@ func TestExecAllocator(t *testing.T) {
 		t.Fatalf("wanted %q, got %q", want, got)
 	}
 
-	// TODO: make cancel() work here too
+	cancel()
-	poolCancel()
 
 	tempDir := FromContext(taskCtx).Browser.userDataDir
 	if _, err := os.Lstat(tempDir); !os.IsNotExist(err) {

context.go

@@ -208,6 +208,7 @@ func Targets(ctx context.Context) ([]*target.Info, error) {
 			return nil, err
 		}
 		c.Browser = browser
+		c.first = true
 	}
 	return target.GetTargets().Do(ctx, c.Browser)
 }