imgbundler: Fixes

- Make bundle flag work
- Display error and update render at the same time in watch mode.
  Before we would just display the render and not show the error.

- Rename imgbundler.InlineX functions to BundleX
- Print imgbundler fetch/readFile errors as they happen in the workers
  instead of coalescing and printing at the end.
- Minor performance improvements by using []byte everywhere possible.
- Improved symbol naming in imgbundler code
- **major**: Ignore already bundled images instead of trying to os.ReadFile them.

cmd/d2/main.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/playwright-community/playwright-go"
 	"github.com/spf13/pflag"
+	"go.uber.org/multierr"
 
 	"oss.terrastruct.com/d2"
 	"oss.terrastruct.com/d2/d2layouts/d2sequence"
@@ -156,6 +157,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 			port:         *portFlag,
 			inputPath:    inputPath,
 			outputPath:   outputPath,
+			bundle:       *bundleFlag,
 			pw:           pw,
 		})
 		if err != nil {
@@ -167,19 +169,15 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	ctx, cancel := context.WithTimeout(ctx, time.Minute*2)
 	defer cancel()
 
-	if *bundleFlag {
+	_, err = compile(ctx, ms, plugin, *themeFlag, inputPath, outputPath, *bundleFlag, pw.Page)
-		_ = 343
-	}
-
-	_, err = compile(ctx, ms, false, plugin, *themeFlag, inputPath, outputPath, pw.Page)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to compile: %w", err)
 	}
 	ms.Log.Success.Printf("successfully compiled %v to %v", inputPath, outputPath)
 	return nil
 }
 
-func compile(ctx context.Context, ms *xmain.State, isWatching bool, plugin d2plugin.Plugin, themeID int64, inputPath, outputPath string, page playwright.Page) ([]byte, error) {
+func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, themeID int64, inputPath, outputPath string, bundle bool, page playwright.Page) ([]byte, error) {
 	input, err := ms.ReadPath(inputPath)
 	if err != nil {
 		return nil, err
@@ -210,38 +208,37 @@ func compile(ctx context.Context, ms *xmain.State, isWatching bool, plugin d2plu
 	}
 	svg, err = plugin.PostProcess(ctx, svg)
 	if err != nil {
-		return nil, err
+		return svg, err
 	}
-	svg, err = imgbundler.InlineLocal(ctx, ms, svg)
+
-	if err != nil {
+	svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg)
-		ms.Log.Error.Printf("missing/broken local image(s), writing partial output: %v", err)
+	if bundle {
+		var bundleErr2 error
+		svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg)
+		bundleErr = multierr.Combine(bundleErr, bundleErr2)
 	}
 
 	out := svg
 	if filepath.Ext(outputPath) == ".png" {
-		svg, err = imgbundler.InlineRemote(ctx, ms, svg)
+		svg := svg
-		if err != nil {
+		if !bundle {
-			ms.Log.Error.Printf("missing/broken remote image(s), writing partial output: %v", err)
+			var bundleErr2 error
+			svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg)
+			bundleErr = multierr.Combine(bundleErr, bundleErr2)
 		}
 
 		out, err = png.ConvertSVG(ms, page, svg)
 		if err != nil {
-			return nil, err
+			return svg, err
 		}
 	}
 
 	err = ms.WritePath(outputPath, out)
 	if err != nil {
-		return nil, err
+		return svg, err
 	}
 
-	// Missing/broken images are fine during watch mode, as the user is likely building up a diagram.
+	return svg, bundleErr
-	// Otherwise, the assumption is that this diagram is building for production, and broken images are not okay.
-	if !isWatching && ms.Log.Nerrors() > 0 {
-		return nil, xmain.ExitErrorf(1, "errors logged while rendering, partial output written to %v", outputPath)
-	}
-
-	return svg, nil
 }
 
 // newExt must include leading .

cmd/d2/static/watch.js

@@ -4,6 +4,9 @@ window.addEventListener("DOMContentLoaded", () => {
 });
 
 function init(reconnectDelay) {
+  const d2ErrDiv = window.document.querySelector("#d2-err");
+  const d2SVG = window.document.querySelector("#d2-svg");
+
   const devMode = document.body.dataset.d2DevMode === "true";
   const ws = new WebSocket(
     `ws://${window.location.host}${window.location.pathname}watch`
@@ -19,13 +22,7 @@ function init(reconnectDelay) {
     } else {
       console.debug("watch websocket received data");
     }
-    const d2ErrDiv = window.document.querySelector("#d2-err");
+    if (msg.svg) {
-    if (msg.err) {
-      d2ErrDiv.innerText = msg.err;
-      d2ErrDiv.style.display = "block";
-      d2ErrDiv.scrollIntoView();
-    } else {
-      const d2SVG = window.document.querySelector("#d2-svg");
       // We could turn d2SVG into an actual SVG element and use outerHTML to fully replace it
       // with the result from the renderer but unfortunately that overwrites the #d2-svg ID.
       // Even if you add another line to set it afterwards. The parsing/interpretation of outerHTML must be async.
@@ -36,6 +33,11 @@ function init(reconnectDelay) {
       d2SVG.innerHTML = msg.svg;
       d2ErrDiv.style.display = "none";
     }
+    if (msg.err) {
+      d2ErrDiv.innerText = msg.err;
+      d2ErrDiv.style.display = "block";
+      d2ErrDiv.scrollIntoView();
+    }
   };
   ws.onerror = (ev) => {
     console.error("watch websocket connection error", ev);

cmd/d2/watch.go

@@ -42,6 +42,7 @@ type watcherOpts struct {
 	port         string
 	inputPath    string
 	outputPath   string
+	bundle       bool
 	pw           png.Playwright
 }
 
@@ -73,8 +74,8 @@ type watcher struct {
 }
 
 type compileResult struct {
-	Err string `json:"err"`
 	SVG string `json:"svg"`
+	Err string `json:"err"`
 }
 
 func newWatcher(ctx context.Context, ms *xmain.State, opts watcherOpts) (*watcher, error) {
@@ -345,19 +346,23 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 			w.pw = newPW
 		}
 
-		b, err := compile(ctx, w.ms, true, w.layoutPlugin, w.themeID, w.inputPath, w.outputPath, w.pw.Page)
+		b, err := compile(ctx, w.ms, w.layoutPlugin, w.themeID, w.inputPath, w.outputPath, w.bundle, w.pw.Page)
+		errs := ""
 		if err != nil {
+			if len(b) > 0 {
+				err = fmt.Errorf("failed to %scompile (rendering partial output): %w", recompiledPrefix, err)
+			} else {
 				err = fmt.Errorf("failed to %scompile: %w", recompiledPrefix, err)
-			w.ms.Log.Error.Print(err)
+			}
-			w.broadcast(&compileResult{
+			errs = err.Error()
-				Err: err.Error(),
+			w.ms.Log.Error.Print(errs)
-			})
 		} else {
 			w.ms.Log.Success.Printf("successfully %scompiled %v to %v", recompiledPrefix, w.inputPath, w.outputPath)
+		}
 		w.broadcast(&compileResult{
 			SVG: string(b),
+			Err: errs,
 		})
-		}
 
 		if firstCompile {
 			firstCompile = false

go.mod

@@ -20,6 +20,7 @@ require (
 	golang.org/x/image v0.1.0
 	golang.org/x/net v0.2.0
 	golang.org/x/text v0.4.0
+	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2
 	gonum.org/v1/plot v0.12.0
 	nhooyr.io/websocket v1.8.7
 	oss.terrastruct.com/cmdlog v0.0.0-20221129200109-540ef52ff07d
@@ -57,7 +58,6 @@ require (
 	golang.org/x/crypto v0.3.0 // indirect
 	golang.org/x/sys v0.2.0 // indirect
 	golang.org/x/term v0.2.0 // indirect
-	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/genproto v0.0.0-20220822174746-9e6da59bd2fc // indirect
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect

lib/imgbundler/imgbundler.go

@@ -14,7 +14,7 @@ import (
 	"sync"
 	"time"
 
-	"go.uber.org/multierr"
+	"golang.org/x/xerrors"
 	"oss.terrastruct.com/xdefer"
 
 	"oss.terrastruct.com/d2/lib/xmain"
@@ -22,140 +22,142 @@ import (
 
 const maxImageSize int64 = 1 << 25 // 33_554_432
 
-var imageRe = regexp.MustCompile(`<image href="([^"]+)"`)
+var imageRegex = regexp.MustCompile(`<image href="([^"]+)"`)
 
-type resp struct {
+func BundleLocal(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
-	srctxt string
+	return bundle(ctx, ms, in, false)
-	data   string
-	err    error
 }
 
-func InlineLocal(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
+func BundleRemote(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
-	return inline(ctx, ms, in, false)
+	return bundle(ctx, ms, in, true)
 }
 
-func InlineRemote(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
+type repl struct {
-	return inline(ctx, ms, in, true)
+	from []byte
+	to   []byte
 }
 
-func inline(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_ []byte, err error) {
+func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_ []byte, err error) {
 	defer xdefer.Errorf(&err, "failed to bundle images")
-	imgs := imageRe.FindAllSubmatch(svg, -1)
+	imgs := imageRegex.FindAllSubmatch(svg, -1)
-
+	imgs = filterImageElements(imgs, isRemote)
-	var filtered [][][]byte
-	for _, img := range imgs {
-		u, err := url.Parse(string(img[1]))
-		isRemoteImg := err == nil && strings.HasPrefix(u.Scheme, "http")
-		if isRemoteImg == isRemote {
-			filtered = append(filtered, img)
-		}
-	}
 
 	var wg sync.WaitGroup
-	respChan := make(chan resp)
+	replc := make(chan repl)
 	// Limits the number of workers to 16.
 	sema := make(chan struct{}, 16)
 
+	var errhrefsMu sync.Mutex
+	var errhrefs []string
+
 	ctx, cancel := context.WithTimeout(ctx, time.Minute*5)
 	defer cancel()
 
-	wg.Add(len(filtered))
+	wg.Add(len(imgs))
 	// Start workers as the sema allows.
 	go func() {
-		for _, img := range filtered {
+		for _, img := range imgs {
 			sema <- struct{}{}
-			go func(src, href string) {
+			go func(imgel, href []byte) {
 				defer func() {
 					wg.Done()
 					<-sema
 				}()
 
-				var data string
+				var buf []byte
 				var err error
 				if isRemote {
-					data, err = fetch(ctx, href)
+					buf, err = httpGet(ctx, string(href))
 				} else {
-					data, err = read(href)
+					buf, err = os.ReadFile(string(href))
 				}
+				if err != nil {
+					ms.Log.Error.Printf("failed to bundle %s: %v", imgel, err)
+					errhrefsMu.Lock()
+					errhrefs = append(errhrefs, string(href))
+					errhrefsMu.Unlock()
+					return
+				}
+
+				mimeType := http.DetectContentType(buf)
+				mimeType = strings.Replace(mimeType, "text/xml", "image/svg+xml", 1)
+				b64 := base64.StdEncoding.EncodeToString(buf)
+
 				select {
 				case <-ctx.Done():
-				case respChan <- resp{
+				case replc <- repl{
-					srctxt: src,
+					from: imgel,
-					data:   data,
+					to:   []byte(fmt.Sprintf(`<image href="data:%s;base64,%s"`, mimeType, b64)),
-					err:    err,
 				}:
 				}
-			}(string(img[0]), string(img[1]))
+			}(img[0], img[1])
 		}
 	}()
 
 	go func() {
 		wg.Wait()
-		close(respChan)
+		close(replc)
 	}()
 
 	for {
 		select {
 		case <-ctx.Done():
-			return nil, fmt.Errorf("failed to wait for imgbundler workers: %w", ctx.Err())
+			return svg, xerrors.Errorf("failed to wait for workers: %w", ctx.Err())
 		case <-time.After(time.Second * 5):
 			ms.Log.Info.Printf("fetching images...")
-		case resp, ok := <-respChan:
+		case repl, ok := <-replc:
 			if !ok {
-				return svg, err
+				if len(errhrefs) > 0 {
+					return svg, xerrors.Errorf("failed to bundle the following images: %v", errhrefs)
 				}
-			if resp.err != nil {
+				return svg, nil
-				err = multierr.Combine(err, resp.err)
+			}
+			svg = bytes.Replace(svg, repl.from, repl.to, 1)
+		}
+	}
+}
+
+// filterImageElements finds all image elements in imgs that are eligible
+// for bundling in the current context.
+func filterImageElements(imgs [][][]byte, isRemote bool) [][][]byte {
+	imgs2 := imgs[:0]
+	for _, img := range imgs {
+		href := string(img[1])
+
+		// Skip already bundled images.
+		if strings.HasPrefix(href, "data:") {
 			continue
 		}
-			svg = bytes.Replace(svg, []byte(resp.srctxt), []byte(fmt.Sprintf(`<image href="%s"`, resp.data)), 1)
+
+		u, err := url.Parse(href)
+		isRemoteImg := err == nil && strings.HasPrefix(u.Scheme, "http")
+
+		if isRemoteImg == isRemote {
+			imgs2 = append(imgs2, img)
 		}
 	}
+	return imgs2
 }
 
-var imgClient = &http.Client{}
+var httpClient = &http.Client{}
 
-func fetch(ctx context.Context, href string) (string, error) {
+func httpGet(ctx context.Context, href string) ([]byte, error) {
 	ctx, cancel := context.WithTimeout(ctx, time.Minute)
 	defer cancel()
 
 	req, err := http.NewRequestWithContext(ctx, "GET", href, nil)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
-	imgResp, err := imgClient.Do(req)
+	resp, err := httpClient.Do(req)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
-	defer imgResp.Body.Close()
+	defer resp.Body.Close()
-	if imgResp.StatusCode != 200 {
+	if resp.StatusCode != 200 {
-		return "", fmt.Errorf("img %s returned status code %d", href, imgResp.StatusCode)
+		return nil, fmt.Errorf("expected status 200 but got %d %s", resp.StatusCode, resp.Status)
 	}
-	r := http.MaxBytesReader(nil, imgResp.Body, maxImageSize)
+	r := http.MaxBytesReader(nil, resp.Body, maxImageSize)
-	data, err := ioutil.ReadAll(r)
+	return ioutil.ReadAll(r)
-	if err != nil {
-		return "", err
-	}
-
-	mimeType := http.DetectContentType(data)
-	mimeType = strings.Replace(mimeType, "text/xml", "image/svg+xml", 1)
-
-	enc := base64.StdEncoding.EncodeToString(data)
-
-	return fmt.Sprintf("data:%s;base64,%s", mimeType, enc), nil
-}
-
-func read(href string) (string, error) {
-	data, err := os.ReadFile(href)
-	if err != nil {
-		return "", err
-	}
-
-	mimeType := http.DetectContentType(data)
-	mimeType = strings.Replace(mimeType, "text/xml", "image/svg+xml", 1)
-
-	enc := base64.StdEncoding.EncodeToString(data)
-
-	return fmt.Sprintf("data:%s;base64,%s", mimeType, enc), nil
 }

lib/imgbundler/imgbundler_test.go

@@ -41,7 +41,7 @@ func TestRegex(t *testing.T) {
 
 	for _, href := range append(urls, notURLs...) {
 		str := fmt.Sprintf(`<image href="%s" />`, href)
-		matches := imageRe.FindAllStringSubmatch(str, -1)
+		matches := imageRegex.FindAllStringSubmatch(str, -1)
 		if len(matches) != 1 {
 			t.Fatalf("uri regex didn't match %s", str)
 		}
@@ -90,7 +90,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 	}
 	ms.Log = cmdlog.Log(ms.Env, os.Stderr)
 
-	imgClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
 		switch req.URL.String() {
 		case svgURL:
@@ -104,7 +104,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		return respRecorder.Result()
 	})
 
-	out, err := InlineRemote(ctx, ms, []byte(sampleSVG))
+	out, err := BundleRemote(ctx, ms, []byte(sampleSVG))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -119,7 +119,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 	}
 
 	// Test almost too large response
-	imgClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
 		bytes := make([]byte, maxImageSize)
 		rand.Read(bytes)
@@ -127,13 +127,13 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		respRecorder.WriteHeader(200)
 		return respRecorder.Result()
 	})
-	_, err = InlineRemote(ctx, ms, []byte(sampleSVG))
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	// Test too large response
-	imgClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
 		bytes := make([]byte, maxImageSize+1)
 		rand.Read(bytes)
@@ -141,18 +141,18 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		respRecorder.WriteHeader(200)
 		return respRecorder.Result()
 	})
-	_, err = InlineRemote(ctx, ms, []byte(sampleSVG))
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
 	if err == nil {
 		t.Fatal("expected error")
 	}
 
 	// Test error response
-	imgClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
 		respRecorder.WriteHeader(500)
 		return respRecorder.Result()
 	})
-	_, err = InlineRemote(ctx, ms, []byte(sampleSVG))
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
 	if err == nil {
 		t.Fatal("expected error")
 	}
@@ -205,7 +205,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		Env: xos.NewEnv(os.Environ()),
 	}
 	ms.Log = cmdlog.Log(ms.Env, os.Stderr)
-	out, err := InlineLocal(ctx, ms, []byte(sampleSVG))
+	out, err := BundleLocal(ctx, ms, []byte(sampleSVG))
 	if err != nil {
 		t.Fatal(err)
 	}