From 687cf6d766a7eaa8d034783919c52bada10c5a38 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Sun, 7 Apr 2019 23:40:29 +0200
Subject: [PATCH] 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      | 9 ++++++++-
 allocate_test.go | 3 +--
 context.go       | 1 +
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/allocate.go b/allocate.go
index ee94f14..390de2f 100644
--- a/allocate.go
+++ b/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
diff --git a/allocate_test.go b/allocate_test.go
index e8d2b5b..7cfe3e6 100644
--- a/allocate_test.go
+++ b/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
-	poolCancel()
+	cancel()
 
 	tempDir := FromContext(taskCtx).Browser.userDataDir
 	if _, err := os.Lstat(tempDir); !os.IsNotExist(err) {
diff --git a/context.go b/context.go
index 1ccb8c6..eae2e9a 100644
--- a/context.go
+++ b/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)
 }