From bdb4f6e324c2fa4500f4cb1d8ccc5e06eb885d92 Mon Sep 17 00:00:00 2001
From: Yad Smood <y.s.inside@gmail.com>
Date: Sun, 30 Jul 2023 00:43:06 +0800
Subject: [PATCH] add AutoFree

---
 cmd/serve/main.go |  3 ++-
 service.go        | 29 ++++++++++++++++++++++++++---
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/cmd/serve/main.go b/cmd/serve/main.go
index de2b309..059d46c 100644
--- a/cmd/serve/main.go
+++ b/cmd/serve/main.go
@@ -37,7 +37,8 @@ func main() {
 	b.BlockRequests(blockList...)
 	b.BypassUserAgentNames(bypassUAs...)
 	b.MaxWait(*maxWait)
-	b.WarnUp()
+	b.WarmUp()
+	b.AutoFree()
 
 	err := http.ListenAndServe(*port, b)
 	if err != nil {
diff --git a/service.go b/service.go
index 52ea658..f3a0500 100644
--- a/service.go
+++ b/service.go
@@ -73,6 +73,10 @@ func (b *Bartender) MaxWait(d time.Duration) {
 	b.maxWait = d
 }
 
+func (b *Bartender) getPage() *rod.Page {
+	return b.pool.Get(b.newPage)
+}
+
 func (b *Bartender) newPage() *rod.Page {
 	page := rod.New().MustConnect().MustPage()
 
@@ -93,12 +97,31 @@ func (b *Bartender) newPage() *rod.Page {
 	return page
 }
 
-func (b *Bartender) WarnUp() {
+// WarmUp pre-creates the headless browsers.
+func (b *Bartender) WarmUp() {
 	for i := 0; i < len(b.pool); i++ {
-		b.pool.Put(b.pool.Get(b.newPage))
+		b.pool.Put(b.getPage())
 	}
 }
 
+// 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() {
+	go func() {
+		for {
+			time.Sleep(10 * time.Minute)
+
+			err := b.getPage().Browser().Close()
+			if err != nil {
+				log.Println("failed to close browser:", err)
+
+				continue
+			}
+			b.pool.Put(nil)
+		}
+	}()
+}
+
 func (b *Bartender) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ua := useragent.Parse(r.Header.Get("User-Agent"))
 	if r.Method != http.MethodGet || b.bypassList[ua.Name] {
@@ -138,7 +161,7 @@ func (b *Bartender) RenderPage(w http.ResponseWriter, r *http.Request) bool {
 
 	w.WriteHeader(statusCode)
 
-	page := b.pool.Get(b.newPage)
+	page := b.getPage()
 	defer b.pool.Put(page)
 
 	page, cancel := page.WithCancel()