create Playwright struct, use BrowserContext

cmd/d2/main.go

@@ -113,15 +113,14 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	}
 	ms.Log.Debug.Printf("using layout plugin %s (%s)", envD2Layout, pluginLocation)
 
-	var pw *playwright.Playwright
+	var pw png.Playwright
-	var browser playwright.Browser
 	if filepath.Ext(outputPath) == ".png" {
-		pw, browser, err = png.InitPlaywright()
+		pw, err = png.InitPlaywright()
 		if err != nil {
 			return err
 		}
 		defer func() error {
-			err = png.Cleanup(pw, browser)
+			err = pw.Cleanup(*watchFlag)
 			if err != nil {
 				return err
 			}
@@ -134,7 +133,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 			return xmain.UsageErrorf("-w[atch] cannot be combined with reading input from stdin")
 		}
 		ms.Env.Setenv("LOG_TIMESTAMPS", "1")
-		w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath, pw, browser)
+		w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath, pw)
 		if err != nil {
 			return err
 		}
@@ -148,7 +147,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 		_ = 343
 	}
 
-	_, err = compile(ctx, ms, plugin, inputPath, outputPath, browser)
+	_, err = compile(ctx, ms, plugin, inputPath, outputPath, pw.Page)
 	if err != nil {
 		return err
 	}
@@ -157,7 +156,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	return nil
 }
 
-func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string, browser playwright.Browser) ([]byte, error) {
+func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string, page playwright.Page) ([]byte, error) {
 	input, err := ms.ReadPath(inputPath)
 	if err != nil {
 		return nil, err
@@ -188,7 +187,7 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, input
 	}
 
 	if filepath.Ext(outputPath) == ".png" {
-		outputImage, err = png.ExportPNG(browser, svg)
+		outputImage, err = png.ExportPNG(ms, page, svg)
 		if err != nil {
 			return nil, err
 		}

cmd/d2/watch.go

@@ -16,11 +16,11 @@ import (
 	"time"
 
 	"github.com/fsnotify/fsnotify"
-	"github.com/playwright-community/playwright-go"
 	"nhooyr.io/websocket"
 	"nhooyr.io/websocket/wsjson"
 
 	"oss.terrastruct.com/d2/d2plugin"
+	"oss.terrastruct.com/d2/lib/png"
 	"oss.terrastruct.com/d2/lib/xbrowser"
 	"oss.terrastruct.com/d2/lib/xhttp"
 	"oss.terrastruct.com/d2/lib/xmain"
@@ -63,8 +63,7 @@ type watcher struct {
 	resMu sync.Mutex
 	res   *compileResult
 
-	browser playwright.Browser
+	pw png.Playwright
-	pw      *playwright.Playwright
 }
 
 type compileResult struct {
@@ -72,7 +71,7 @@ type compileResult struct {
 	SVG string `json:"svg"`
 }
 
-func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string, pw *playwright.Playwright, browser playwright.Browser) (*watcher, error) {
+func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string, pw png.Playwright) (*watcher, error) {
 	ctx, cancel := context.WithCancel(ctx)
 
 	w := &watcher{
@@ -87,7 +86,6 @@ func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plug
 
 		compileCh: make(chan struct{}, 1),
 		wsclients: make(map[*wsclient]struct{}),
-		browser:   browser,
 		pw:        pw,
 	}
 	err := w.init()
@@ -332,15 +330,16 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 			recompiledPrefix = "re"
 		}
 
-		if !w.browser.IsConnected() {
+		if filepath.Ext(w.outputPath) == ".png" && !w.pw.Browser.IsConnected() {
-			newBrowser, err := w.pw.Chromium.Launch()
+			newPW, err := w.pw.RestartBrowser()
 			if err != nil {
+				w.ms.Log.Error.Printf("failed to refresh Playwright browser")
 				return err
 			}
-			w.browser = newBrowser
+			w.pw = newPW
 		}
 
-		b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath, w.browser)
+		b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath, w.pw.Page)
 		if err != nil {
 			err = fmt.Errorf("failed to %scompile: %w", recompiledPrefix, err)
 			w.ms.Log.Error.Print(err)

lib/png/png.go

@@ -12,19 +12,69 @@ import (
 	_ "embed"
 
 	"github.com/playwright-community/playwright-go"
+	"oss.terrastruct.com/d2/lib/xmain"
 )
 
-func InitPlaywright() (*playwright.Playwright, playwright.Browser, error) {
+type Playwright struct {
+	PW             *playwright.Playwright
+	Browser        playwright.Browser
+	BrowserContext playwright.BrowserContext
+	Page           playwright.Page
+}
+
+func (pw *Playwright) RestartBrowser() (newPW Playwright, err error) {
+	if err = pw.BrowserContext.Close(); err != nil {
+		return Playwright{}, err
+	}
+	if err = pw.Browser.Close(); err != nil {
+		return Playwright{}, err
+	}
+	browser, err := pw.PW.Chromium.Launch()
+	if err != nil {
+		return Playwright{}, err
+	}
+	context, err := browser.NewContext()
+	if err != nil {
+		return Playwright{}, err
+	}
+	page, err := context.NewPage()
+	if err != nil {
+		return Playwright{}, err
+	}
+	return Playwright{
+		PW:             pw.PW,
+		Browser:        browser,
+		BrowserContext: context,
+		Page:           page,
+	}, nil
+}
+
+func (pw *Playwright) Cleanup(isWatch bool) (err error) {
+	if !isWatch {
+		if err = pw.BrowserContext.Close(); err != nil {
+			return err
+		}
+	}
+	if err = pw.Browser.Close(); err != nil {
+		return err
+	}
+	if err = pw.PW.Stop(); err != nil {
+		return err
+	}
+	return nil
+}
+
+func InitPlaywright() (Playwright, error) {
 	// check if playwright driver/browsers are installed and up to date
 	// https://github.com/playwright-community/playwright-go/blob/8e8f670b5fa7ba5365ae4bfc123fea4aac359763/run.go#L64.
 	driver, err := playwright.NewDriver(&playwright.RunOptions{})
 	if err != nil {
-		return nil, nil, err
+		return Playwright{}, err
 	}
 	if _, err := os.Stat(driver.DriverBinaryLocation); errors.Is(err, os.ErrNotExist) {
 		err = playwright.Install()
 		if err != nil {
-			return nil, nil, err
+			return Playwright{}, err
 		}
 	} else if err == nil {
 		cmd := exec.Command(driver.DriverBinaryLocation, "--version")
@@ -32,42 +82,44 @@ func InitPlaywright() (*playwright.Playwright, playwright.Browser, error) {
 		if err != nil || !bytes.Contains(output, []byte(driver.Version)) {
 			err = playwright.Install()
 			if err != nil {
-				return nil, nil, err
+				return Playwright{}, err
 			}
 		}
 	}
 
 	pw, err := playwright.Run()
 	if err != nil {
-		return nil, nil, err
+		return Playwright{}, err
 	}
 	browser, err := pw.Chromium.Launch()
 	if err != nil {
-		return nil, nil, err
+		return Playwright{}, err
 	}
-	return pw, browser, nil
+	context, err := browser.NewContext()
+	if err != nil {
+		return Playwright{}, err
+	}
+	page, err := context.NewPage()
+	if err != nil {
+		return Playwright{}, err
+	}
+	return Playwright{
+		PW:             pw,
+		Browser:        browser,
+		BrowserContext: context,
+		Page:           page,
+	}, nil
 }
 
 //go:embed generate_png.js
 var genPNGScript string
 
-func ExportPNG(browser playwright.Browser, svg []byte) (outputImage []byte, err error) {
+func ExportPNG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage []byte, err error) {
-	var page playwright.Page
+	if page == nil {
-	defer func() error {
+		ms.Log.Error.Printf("Playwright was not initialized properly for PNG export")
-		err = page.Close()
+		return nil, fmt.Errorf("Playwright page is not initialized for png export")
-		if err != nil {
+	}
-			return err
-		}
-		return nil
-	}()
 
-	if browser == nil {
-		return nil, fmt.Errorf("browser is not initialized for png export")
-	}
-	page, err = browser.NewPage()
-	if err != nil {
-		return nil, err
-	}
 	encodedSVG := base64.StdEncoding.EncodeToString(svg)
 	pngInterface, err := page.Evaluate(genPNGScript, "data:image/svg+xml;charset=utf-8;base64,"+encodedSVG)
 	if err != nil {
@@ -77,27 +129,9 @@ func ExportPNG(browser playwright.Browser, svg []byte) (outputImage []byte, err
 	pngString := fmt.Sprintf("%v", pngInterface)
 	pngPrefix := "data:image/png;base64,"
 	if !strings.HasPrefix(pngString, pngPrefix) {
+		ms.Log.Error.Printf("failed to convert D2 file to PNG")
 		return nil, fmt.Errorf("playwright export generated invalid png")
 	}
 	splicedPNGString := pngString[len(pngPrefix):]
-	outputImage, err = base64.StdEncoding.DecodeString(splicedPNGString)
+	return base64.StdEncoding.DecodeString(splicedPNGString)
-	if err != nil {
-		return nil, err
-	}
-
-	return outputImage, nil
-}
-
-func Cleanup(pw *playwright.Playwright, browser playwright.Browser) (err error) {
-	if browser != nil {
-		if err = browser.Close(); err != nil {
-			return err
-		}
-	}
-	if pw != nil {
-		if err = pw.Stop(); err != nil {
-			return err
-		}
-	}
-	return nil
 }