diff --git a/cmd/d2/main.go b/cmd/d2/main.go index 6331db81c..8f86b270d 100644 --- a/cmd/d2/main.go +++ b/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 browser playwright.Browser + var pw png.Playwright 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 } diff --git a/cmd/d2/watch.go b/cmd/d2/watch.go index d1f31056b..0930caf5b 100644 --- a/cmd/d2/watch.go +++ b/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 *playwright.Playwright + pw png.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() { - newBrowser, err := w.pw.Chromium.Launch() + if filepath.Ext(w.outputPath) == ".png" && !w.pw.Browser.IsConnected() { + 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) diff --git a/lib/png/png.go b/lib/png/png.go index a32d74179..d3c0f4a74 100644 --- a/lib/png/png.go +++ b/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) { - var page playwright.Page - defer func() error { - err = page.Close() - if err != nil { - return err - } - return nil - }() +func ExportPNG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage []byte, err error) { + if page == nil { + ms.Log.Error.Printf("Playwright was not initialized properly for PNG export") + return nil, fmt.Errorf("Playwright page is not initialized for png export") + } - 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) - 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 + return base64.StdEncoding.DecodeString(splicedPNGString) }