also release files

Author: ysmood

cmd/serve/main.go

@@ -15,6 +15,7 @@ func main() {
 	target := flag.String("t", "", "target url to proxy")
 	size := flag.Int("s", 2, "size of the pool")
 	maxWait := flag.Duration("w", 3*time.Second, "max wait time for a page rendering")
+	autoFree := flag.Duration("f", 10*time.Minute, "auto close each headless browser after the specified time")
 
 	var bypassUAs StringsFlag = bartender.DefaultBypassUserAgentNames
 	flag.Var(&bypassUAs, "u", "bypass the specified user-agent names")
@@ -38,7 +39,7 @@ func main() {
 	b.BypassUserAgentNames(bypassUAs...)
 	b.MaxWait(*maxWait)
 	b.WarmUp()
-	b.AutoFree()
+	b.AutoFree(*autoFree)
 
 	err := http.ListenAndServe(*port, b)
 	if err != nil {

service.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/go-rod/rod"
+	"github.com/go-rod/rod/lib/launcher"
 	"github.com/go-rod/rod/lib/proto"
 	"github.com/mileusna/useragent"
 )
@@ -78,7 +79,10 @@ func (b *Bartender) getPage() *rod.Page {
 }
 
 func (b *Bartender) newPage() *rod.Page {
-	page := rod.New().MustConnect().MustPage()
+	l := launcher.New()
+	go l.Cleanup()
+
+	page := rod.New().ControlURL(l.MustLaunch()).MustConnect().MustPage()
 
 	if len(b.blockRequests) > 0 {
 		router := page.HijackRequests()
@@ -106,18 +110,22 @@ func (b *Bartender) WarmUp() {
 
 // AutoFree automatically closes the each headless browser after a period of time.
 // It prevent the memory leak of the headless browser.
-func (b *Bartender) AutoFree() {
+func (b *Bartender) AutoFree(interval time.Duration) {
 	go func() {
 		for {
-			time.Sleep(10 * time.Minute)
+			time.Sleep(interval)
+
+			page := b.getPage()
+			browser := page.Browser()
 
-			err := b.getPage().Browser().Close()
+			err := browser.Close()
 			if err != nil {
 				log.Println("failed to close browser:", err)
 
 				continue
 			}
 			b.pool.Put(nil)
+			log.Println("headless browser freed:", page.SessionID)
 		}
 	}()
 }