From d0484ed1c5e74bf82f3f1d5aaf4c71315a5d3bf9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Mon, 8 Apr 2019 12:10:59 +0200
Subject: [PATCH] simplify the allocator API

Exposing NewAllocator and AllocatorOption was unnecessary, and it made
the API more complex to use and understand.

Instead, have users call NewExecAllocator directly. This removes some
code, and simplifies the examples and tests.
---
 allocate.go      | 46 +++++++++++++++++++---------------------------
 allocate_test.go |  4 ++--
 chromedp_test.go |  2 +-
 context.go       |  8 ++++----
 example_test.go  |  3 +--
 5 files changed, 27 insertions(+), 36 deletions(-)

diff --git a/allocate.go b/allocate.go
index 390de2f..66bf79f 100644
--- a/allocate.go
+++ b/allocate.go
@@ -31,15 +31,26 @@ type Allocator interface {
 	Wait()
 }
 
-// NewAllocator creates a new allocator context, suitable for use with
-// NewContext or Run.
-func NewAllocator(parent context.Context, opts ...AllocatorOption) (context.Context, context.CancelFunc) {
-	ctx, cancel := context.WithCancel(parent)
-	c := &Context{}
-
-	for _, o := range opts {
-		o(&c.Allocator)
+// setupExecAllocator is similar to NewExecAllocator, but it allows NewContext
+// to create the allocator without the unnecessary context layer.
+func setupExecAllocator(opts ...ExecAllocatorOption) *ExecAllocator {
+	ep := &ExecAllocator{
+		initFlags: make(map[string]interface{}),
 	}
+	for _, o := range opts {
+		o(ep)
+	}
+	if ep.execPath == "" {
+		ep.execPath = findExecPath()
+	}
+	return ep
+}
+
+// NewExecAllocator creates a new context set up with an ExecAllocator, suitable
+// for use with NewContext or Run.
+func NewExecAllocator(parent context.Context, opts ...ExecAllocatorOption) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(parent)
+	c := &Context{Allocator: setupExecAllocator(opts...)}
 
 	ctx = context.WithValue(ctx, contextKey{}, c)
 	cancelWait := func() {
@@ -49,25 +60,6 @@ func NewAllocator(parent context.Context, opts ...AllocatorOption) (context.Cont
 	return ctx, cancelWait
 }
 
-// AllocatorOption is a allocator option.
-type AllocatorOption func(*Allocator)
-
-// WithExecAllocator returns an AllocatorOption which sets up an ExecAllocator.
-func WithExecAllocator(opts ...ExecAllocatorOption) func(*Allocator) {
-	return func(p *Allocator) {
-		ep := &ExecAllocator{
-			initFlags: make(map[string]interface{}),
-		}
-		for _, o := range opts {
-			o(ep)
-		}
-		if ep.execPath == "" {
-			ep.execPath = findExecPath()
-		}
-		*p = ep
-	}
-}
-
 // ExecAllocatorOption is a exec allocator option.
 type ExecAllocatorOption func(*ExecAllocator)
 
diff --git a/allocate_test.go b/allocate_test.go
index 7cfe3e6..af1c9f5 100644
--- a/allocate_test.go
+++ b/allocate_test.go
@@ -9,7 +9,7 @@ import (
 func TestExecAllocator(t *testing.T) {
 	t.Parallel()
 
-	poolCtx, poolCancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
+	poolCtx, poolCancel := NewExecAllocator(context.Background(), allocOpts...)
 	defer poolCancel()
 
 	// TODO: test that multiple child contexts are run in different
@@ -41,7 +41,7 @@ func TestExecAllocator(t *testing.T) {
 func TestExecAllocatorCancelParent(t *testing.T) {
 	t.Parallel()
 
-	poolCtx, poolCancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
+	poolCtx, poolCancel := NewExecAllocator(context.Background(), allocOpts...)
 	defer poolCancel()
 
 	// TODO: test that multiple child contexts are run in different
diff --git a/chromedp_test.go b/chromedp_test.go
index eec2d5d..57e2f63 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -53,7 +53,7 @@ func TestMain(m *testing.M) {
 		allocOpts = append(allocOpts, NoSandbox)
 	}
 
-	allocCtx, cancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
+	allocCtx, cancel := NewExecAllocator(context.Background(), allocOpts...)
 
 	// start the browser
 	browserCtx, _ = NewContext(allocCtx)
diff --git a/context.go b/context.go
index eae2e9a..bfcce48 100644
--- a/context.go
+++ b/context.go
@@ -62,11 +62,11 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
 		o(c)
 	}
 	if c.Allocator == nil {
-		WithExecAllocator(
+		c.Allocator = setupExecAllocator(
 			NoFirstRun,
 			NoDefaultBrowserCheck,
 			Headless,
-		)(&c.Allocator)
+		)
 	}
 
 	ctx = context.WithValue(ctx, contextKey{}, c)
@@ -113,8 +113,8 @@ func FromContext(ctx context.Context) *Context {
 }
 
 // Run runs an action against the provided context. The provided context must
-// contain a valid Allocator; typically, that will be created via NewContext or
-// NewAllocator.
+// contain a valid Allocator; typically, that will be created via NewContext, or
+// via one of the allocator constructors like NewExecAllocator.
 func Run(ctx context.Context, actions ...Action) error {
 	c := FromContext(ctx)
 	if c == nil || c.Allocator == nil {
diff --git a/example_test.go b/example_test.go
index cab7315..64be637 100644
--- a/example_test.go
+++ b/example_test.go
@@ -45,8 +45,7 @@ func ExampleExecAllocator() {
 		chromedp.UserDataDir(dir),
 	}
 
-	allocCtx, cancel := chromedp.NewAllocator(context.Background(),
-		chromedp.WithExecAllocator(opts...))
+	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
 	defer cancel()
 
 	taskCtx, cancel := chromedp.NewContext(allocCtx)