Skip to content

Commit

Permalink
add AutoFree
Browse files Browse the repository at this point in the history
  • Loading branch information
ysmood committed Jul 30, 2023
1 parent 912dd97 commit bdb4f6e
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 4 deletions.
3 changes: 2 additions & 1 deletion cmd/serve/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
29 changes: 26 additions & 3 deletions service.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand All @@ -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] {
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit bdb4f6e

Please sign in to comment.