v0.0.3 release

* save memory by cleaning up and resetting everything when all tabs are closed
* split test cases

chromedpproxy.go

@@ -2,6 +2,7 @@ package chromedpproxy
 
 import (
 	"context"
+	"errors"
 	"github.com/chromedp/cdproto/page"
 	"github.com/chromedp/cdproto/target"
 	"github.com/chromedp/chromedp"
@@ -10,17 +11,19 @@ import (
 )
 
 var mutex = sync.RWMutex{}
-var loaded = make(chan bool, 1)
 
-var Context context.Context
+var loaded = make(chan bool, 1)
+var mainContext context.Context
+var mainCancel chan bool
+var totalTargets = 0
 
 // PrepareProxy abstracts chromedp.NewExecAllocator to the use case of this package
 // it accepts listen addresses for both Chrome remote debugging and frontend as configuration
 // it is also a variadic function that accepts extra chromedp.ExecAllocatorOption to be passed to the chromedp.NewExecAllocator
 func PrepareProxy(chromeListenAddr string, frontendListenAddr string, customOpts ...chromedp.ExecAllocatorOption) {
 	// ensure only exactly one context is prepared
 	mutex.Lock()
-	if Context != nil {
+	if mainContext != nil {
 		mutex.Unlock()
 		return
 	}
@@ -45,35 +48,38 @@ func PrepareProxy(chromeListenAddr string, frontendListenAddr string, customOpts
 	go func() {
 		ctx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
 		defer cancel()
-		Context = ctx
+		mainContext = ctx
 		loaded <- true
 		mutex.Unlock()
-		startFrontEnd(frontendListenAddr, chromeListenAddrSplit[1])
+		mainCancel = make(chan bool, 1)
+		defer close(mainCancel)
+		startFrontEnd(frontendListenAddr, chromeListenAddrSplit[1], mainCancel)
 	}()
 }
 
 // NewTab abstracts creating a new tab in the root context
 // it returns a target ID or error
 func NewTab(url string, customOpts ...chromedp.ContextOption) (target.ID, error) {
 	// if context is not prepared, create with default values
-	if Context == nil {
+	if mainContext == nil {
 		PrepareProxy(":9222", ":9221")
 		<-loaded
 	}
 	mutex.Lock()
 	defer mutex.Unlock()
 
 	// create new tab and navigate to URL
-	Context, _ = chromedp.NewContext(Context, customOpts...)
-	err := chromedp.Run(Context, chromedp.Tasks{
+	mainContext, _ = chromedp.NewContext(mainContext, customOpts...)
+	err := chromedp.Run(mainContext, chromedp.Tasks{
 		chromedp.Navigate(url),
 	})
 	if err != nil {
 		return "", err
 	}
+	totalTargets++
 
 	// return target ID
-	chromeContext := chromedp.FromContext(Context)
+	chromeContext := chromedp.FromContext(mainContext)
 	return chromeContext.Target.TargetID, nil
 }
 
@@ -83,20 +89,32 @@ func GetTarget(id target.ID) context.Context {
 	defer mutex.RUnlock()
 
 	// return context from target ID
-	ctx, _ := chromedp.NewContext(Context, chromedp.WithTargetID(id))
+	ctx, _ := chromedp.NewContext(mainContext, chromedp.WithTargetID(id))
 	return ctx
 }
 
 // CloseTarget closes a target by closing the page
+// if the last page has been closed, clean up everything
 // it returns an error if any
 func CloseTarget(id target.ID) error {
 	mutex.Lock()
 	defer mutex.Unlock()
 
-	ctx, cancel := chromedp.NewContext(Context, chromedp.WithTargetID(id))
+	if mainContext == nil {
+		return errors.New("context not prepared or already closed")
+	}
+
+	ctx, cancel := chromedp.NewContext(mainContext, chromedp.WithTargetID(id))
 	defer cancel()
 	if err := chromedp.Run(ctx, page.Close()); err != nil {
 		return err
 	}
+	totalTargets--
+	if totalTargets <= 0 {
+		loaded = make(chan bool, 1)
+		mainContext = nil
+		mainCancel <- true
+		totalTargets = 0
+	}
 	return nil
 }

chromedpproxy_test.go

@@ -7,11 +7,26 @@ import (
 	"testing"
 )
 
-func TestChromeDPProxy(t *testing.T) {
-	// prepare proxy
+func TestClosingTwoTabsAndCreateOneTab(t *testing.T) {
 	PrepareProxy(":9222", ":9221")
+	target1ID, err := NewTab("about:blank")
+	if err != nil {
+		t.Error(err)
+	}
+	target2ID, err := NewTab("about:blank")
+	if err != nil {
+		t.Error(err)
+	}
+	err = CloseTarget(target1ID)
+	if err != nil {
+		t.Error(err)
+	}
+	err = CloseTarget(target2ID)
+	if err != nil {
+		t.Error(err)
+	}
 
-	// test creating proxy, then deleting, then creating again
+	PrepareProxy(":9222", ":9221")
 	targetID, err := NewTab("about:blank")
 	if err != nil {
 		t.Error(err)
@@ -20,15 +35,27 @@ func TestChromeDPProxy(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-
+}
+func TestDoubleClose(t *testing.T) {
+	targetID, err := NewTab("about:blank")
+	if err != nil {
+		t.Error(err)
+	}
 	err = CloseTarget(targetID)
-	t.Logf("expected error: %s", err) // should error if attempting to close again
-	targetID, err = NewTab("https://github.com/jarylc/go-chromedpproxy")
+	if err != nil {
+		t.Error(err)
+	}
+	err = CloseTarget(targetID)
+	if err == nil {
+		t.Error("expected error for attempting to close again not returned")
+	}
+}
+func TestRegularChromeDPUsage(t *testing.T) {
+	targetID, err := NewTab("https://github.com/jarylc/go-chromedpproxy")
 	if err != nil {
 		t.Error(err)
 	}
 
-	// test regular chromedp usage
 	ctx := GetTarget(targetID)
 	result := ""
 	err = chromedp.Run(ctx, chromedp.Tasks{
@@ -41,7 +68,6 @@ func TestChromeDPProxy(t *testing.T) {
 		t.Errorf("expected result to be 'go-chromedpproxy', got %s", result)
 	}
 
-	// cleanup
 	err = CloseTarget(targetID)
 	if err != nil {
 		t.Error(err)

fiber_frontend.go

@@ -246,16 +246,21 @@ const html = `<!DOCTYPE html>
 </html>`
 
 // startFrontEnd starts a blocking Fiber web server that serves the front-end alongside the websocket proxy
-func startFrontEnd(frontendListenAddr string, cdpPort string) {
+func startFrontEnd(frontendListenAddr string, cdpPort string, cancelChan chan bool) {
 	app := fiber.New(fiber.Config{
 		ReduceMemoryUsage:     true,
 		DisableStartupMessage: true,
 	})
 	interrupt := make(chan os.Signal, 1)
+	defer close(interrupt)
 	signal.Notify(interrupt, os.Interrupt)
 	go func() {
-		<-interrupt
-		_ = app.Shutdown()
+		select {
+		case <-interrupt:
+			app.Shutdown()
+		case <-cancelChan:
+			app.Shutdown()
+		}
 	}()
 	app.Get("/", func(c *fiber.Ctx) error {
 		c.Set(fiber.HeaderContentType, fiber.MIMETextHTML)