Add context support to Playwright browser initialization
Refactored Playwright initialization to ensure context propagation. Updated `NewPlayWrightBrowser` and related methods to accept `context.Context` for better cancellation and timeout handling. Improved error resilience and concurrency during browser setup.
This commit is contained in:
@@ -43,7 +43,7 @@ var Flags = BrowserFlags{
|
||||
},
|
||||
}
|
||||
|
||||
func FromCommand(_ context.Context, cmd *cli.Command) (extractor.Browser, error) {
|
||||
func FromCommand(ctx context.Context, cmd *cli.Command) (extractor.Browser, error) {
|
||||
var opts extractor.PlayWrightBrowserOptions
|
||||
|
||||
if ua := cmd.String("user-agent"); ua != "" {
|
||||
@@ -72,5 +72,5 @@ func FromCommand(_ context.Context, cmd *cli.Command) (extractor.Browser, error)
|
||||
|
||||
opts.ShowBrowser = cmd.Bool("visible")
|
||||
|
||||
return extractor.NewPlayWrightBrowser(opts)
|
||||
return extractor.NewPlayWrightBrowser(ctx, opts)
|
||||
}
|
||||
|
||||
269
playwright.go
269
playwright.go
@@ -90,7 +90,7 @@ func playwrightCookieToCookie(cookie playwright.Cookie) Cookie {
|
||||
}
|
||||
}
|
||||
|
||||
func NewPlayWrightBrowser(opts ...PlayWrightBrowserOptions) (Browser, error) {
|
||||
func NewPlayWrightBrowser(ctx context.Context, opts ...PlayWrightBrowserOptions) (Browser, error) {
|
||||
var thirtySeconds = 30 * time.Second
|
||||
opt := PlayWrightBrowserOptions{
|
||||
UserAgent: "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:142.0) Gecko/20100101 Firefox/142.0",
|
||||
@@ -131,126 +131,161 @@ func NewPlayWrightBrowser(opts ...PlayWrightBrowserOptions) (Browser, error) {
|
||||
opt.ShowBrowser = o.ShowBrowser
|
||||
}
|
||||
|
||||
pw, err := playwright.Run()
|
||||
|
||||
if err != nil {
|
||||
err = playwright.Install()
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pw, err = playwright.Run()
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var bt playwright.BrowserType
|
||||
|
||||
switch opt.Browser {
|
||||
case PlayWrightBrowserSelectionChromium:
|
||||
bt = pw.Chromium
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_CHROMIUM")
|
||||
}
|
||||
|
||||
case PlayWrightBrowserSelectionFirefox:
|
||||
bt = pw.Firefox
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_FIREFOX")
|
||||
}
|
||||
|
||||
case PlayWrightBrowserSelectionWebKit:
|
||||
bt = pw.WebKit
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_WEBKIT")
|
||||
}
|
||||
|
||||
default:
|
||||
return nil, ErrInvalidBrowserSelection
|
||||
}
|
||||
var browser playwright.Browser
|
||||
|
||||
var launch = true
|
||||
if opt.PlayWrightServerAddress != "" && !opt.UseLocalOnly {
|
||||
launch = false
|
||||
slog.Info("connecting to playwright server", "address", opt.PlayWrightServerAddress)
|
||||
var timeout float64 = 30000
|
||||
browser, err = bt.Connect(opt.PlayWrightServerAddress, playwright.BrowserTypeConnectOptions{Timeout: &timeout})
|
||||
|
||||
if err != nil {
|
||||
if opt.DontLaunchOnConnectFailure {
|
||||
return nil, err
|
||||
}
|
||||
slog.Warn("failed to connect to playwright server, launching local browser", "err", err)
|
||||
launch = true
|
||||
}
|
||||
}
|
||||
|
||||
if launch {
|
||||
browser, err = bt.Launch(playwright.BrowserTypeLaunchOptions{
|
||||
Headless: playwright.Bool(!opt.ShowBrowser),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var viewport *playwright.Size
|
||||
if opt.Dimensions.Width > 0 && opt.Dimensions.Height > 0 {
|
||||
viewport = &playwright.Size{
|
||||
Width: opt.Dimensions.Width,
|
||||
Height: opt.Dimensions.Height,
|
||||
}
|
||||
}
|
||||
|
||||
var scheme *playwright.ColorScheme
|
||||
|
||||
if opt.DarkMode {
|
||||
scheme = playwright.ColorSchemeDark
|
||||
} else {
|
||||
scheme = playwright.ColorSchemeNoPreference
|
||||
}
|
||||
|
||||
c, err := browser.NewContext(playwright.BrowserNewContextOptions{
|
||||
UserAgent: playwright.String(opt.UserAgent),
|
||||
Viewport: viewport,
|
||||
ColorScheme: scheme,
|
||||
})
|
||||
if err != nil {
|
||||
// Check if context is already done
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if opt.CookieJar != nil {
|
||||
cookies, err := opt.CookieJar.GetAll()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting cookies from cookie jar: %w", err)
|
||||
}
|
||||
|
||||
pwCookies := make([]playwright.OptionalCookie, len(cookies))
|
||||
|
||||
for i, cookie := range cookies {
|
||||
pwCookies[i] = cookieToPlaywrightOptionalCookie(cookie)
|
||||
}
|
||||
|
||||
err = c.AddCookies(pwCookies)
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error adding cookies to browser: %w", err)
|
||||
}
|
||||
type browserResult struct {
|
||||
browser Browser
|
||||
err error
|
||||
}
|
||||
|
||||
return playWrightBrowser{
|
||||
pw: pw,
|
||||
browser: browser,
|
||||
userAgent: opt.UserAgent,
|
||||
timeout: *opt.Timeout,
|
||||
cookieJar: opt.CookieJar,
|
||||
ctx: c,
|
||||
serverAddr: opt.PlayWrightServerAddress,
|
||||
}, nil
|
||||
// Create a channel for the result
|
||||
resultCh := make(chan browserResult, 1)
|
||||
|
||||
// Launch browser initialization in a separate goroutine
|
||||
go func() {
|
||||
pw, err := playwright.Run()
|
||||
|
||||
if err != nil {
|
||||
err = playwright.Install()
|
||||
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, err}
|
||||
return
|
||||
}
|
||||
|
||||
pw, err = playwright.Run()
|
||||
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, err}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var bt playwright.BrowserType
|
||||
|
||||
switch opt.Browser {
|
||||
case PlayWrightBrowserSelectionChromium:
|
||||
bt = pw.Chromium
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_CHROMIUM")
|
||||
}
|
||||
|
||||
case PlayWrightBrowserSelectionFirefox:
|
||||
bt = pw.Firefox
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_FIREFOX")
|
||||
}
|
||||
|
||||
case PlayWrightBrowserSelectionWebKit:
|
||||
bt = pw.WebKit
|
||||
if opt.PlayWrightServerAddress == "" {
|
||||
opt.PlayWrightServerAddress = os.Getenv("PLAYWRIGHT_SERVER_ADDRESS_WEBKIT")
|
||||
}
|
||||
|
||||
default:
|
||||
resultCh <- browserResult{nil, ErrInvalidBrowserSelection}
|
||||
return
|
||||
}
|
||||
var browser playwright.Browser
|
||||
|
||||
var launch = true
|
||||
if opt.PlayWrightServerAddress != "" && !opt.UseLocalOnly {
|
||||
launch = false
|
||||
slog.Info("connecting to playwright server", "address", opt.PlayWrightServerAddress)
|
||||
var timeout float64 = 30000
|
||||
browser, err = bt.Connect(opt.PlayWrightServerAddress, playwright.BrowserTypeConnectOptions{Timeout: &timeout})
|
||||
|
||||
if err != nil {
|
||||
if opt.DontLaunchOnConnectFailure {
|
||||
resultCh <- browserResult{nil, err}
|
||||
return
|
||||
}
|
||||
slog.Warn("failed to connect to playwright server, launching local browser", "err", err)
|
||||
launch = true
|
||||
}
|
||||
}
|
||||
|
||||
if launch {
|
||||
browser, err = bt.Launch(playwright.BrowserTypeLaunchOptions{
|
||||
Headless: playwright.Bool(!opt.ShowBrowser),
|
||||
})
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, err}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var viewport *playwright.Size
|
||||
if opt.Dimensions.Width > 0 && opt.Dimensions.Height > 0 {
|
||||
viewport = &playwright.Size{
|
||||
Width: opt.Dimensions.Width,
|
||||
Height: opt.Dimensions.Height,
|
||||
}
|
||||
}
|
||||
|
||||
var scheme *playwright.ColorScheme
|
||||
|
||||
if opt.DarkMode {
|
||||
scheme = playwright.ColorSchemeDark
|
||||
} else {
|
||||
scheme = playwright.ColorSchemeNoPreference
|
||||
}
|
||||
|
||||
c, err := browser.NewContext(playwright.BrowserNewContextOptions{
|
||||
UserAgent: playwright.String(opt.UserAgent),
|
||||
Viewport: viewport,
|
||||
ColorScheme: scheme,
|
||||
})
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, err}
|
||||
return
|
||||
}
|
||||
|
||||
if opt.CookieJar != nil {
|
||||
cookies, err := opt.CookieJar.GetAll()
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, fmt.Errorf("error getting cookies from cookie jar: %w", err)}
|
||||
return
|
||||
}
|
||||
|
||||
pwCookies := make([]playwright.OptionalCookie, len(cookies))
|
||||
|
||||
for i, cookie := range cookies {
|
||||
pwCookies[i] = cookieToPlaywrightOptionalCookie(cookie)
|
||||
}
|
||||
|
||||
err = c.AddCookies(pwCookies)
|
||||
|
||||
if err != nil {
|
||||
resultCh <- browserResult{nil, fmt.Errorf("error adding cookies to browser: %w", err)}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
resultCh <- browserResult{
|
||||
browser: playWrightBrowser{
|
||||
pw: pw,
|
||||
browser: browser,
|
||||
userAgent: opt.UserAgent,
|
||||
timeout: *opt.Timeout,
|
||||
cookieJar: opt.CookieJar,
|
||||
ctx: c,
|
||||
serverAddr: opt.PlayWrightServerAddress,
|
||||
},
|
||||
err: nil,
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for either context cancellation or browser initialization completion
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case result := <-resultCh:
|
||||
return result.browser, result.err
|
||||
}
|
||||
}
|
||||
|
||||
func (b playWrightBrowser) updateCookies(_ context.Context, page playwright.Page) error {
|
||||
@@ -337,7 +372,7 @@ func deferClose(cl io.Closer) {
|
||||
}
|
||||
|
||||
func Screenshot(ctx context.Context, target string, timeout time.Duration) ([]byte, error) {
|
||||
browser, err := NewPlayWrightBrowser(PlayWrightBrowserOptions{
|
||||
browser, err := NewPlayWrightBrowser(ctx, PlayWrightBrowserOptions{
|
||||
Timeout: &timeout,
|
||||
})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user