cli img fetching optimizations

ci/release/changelogs/next.md

@@ -11,6 +11,8 @@
 - `sql_table` now alternatively takes an array of constraints instead of being limited to a single one. Thanks @satoqz ! [#1245](https://github.com/terrastruct/d2/pull/1245)
 - Constraints in `sql_table` render even if they have no matching abbreviation [#1372](https://github.com/terrastruct/d2/pull/1372)
 - Constraints in `sql_table` sheds their excessive letter-spacing and is padded from the end consistently [#1372](https://github.com/terrastruct/d2/pull/1372)
+- Duplicate image URLs in icons are only fetched once [#1373](https://github.com/terrastruct/d2/pull/1373)
+- In watch mode, images are cached by default across compiles. Can be disabled with flag `--img-cache=0`. [#1373](https://github.com/terrastruct/d2/pull/1373)
 
 #### Bugfixes ⛑️
 

ci/release/template/man/d2.1

@@ -109,6 +109,9 @@ An appendix for tooltips and links is added to PNG exports since they are not in
 .Ns .
 .It Fl d , -debug
 Print debug logs.
+.It Fl -img-cache Ar true
+In watch mode, images used in icons are cached for subsequent compilations. This should be disabled if images might change
+.Ns .
 .It Fl h , -help
 Print usage information and exit.
 .It Fl v , -version

d2cli/main.go

@@ -67,6 +67,10 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 	if err != nil {
 		return err
 	}
+	imgCacheFlag, err := ms.Opts.Bool("IMG_CACHE", "img-cache", "", true, "in watch mode, images used in icons are cached for subsequent compilations. This should be disabled if images might change.")
+	if err != nil {
+		return err
+	}
 	layoutFlag := ms.Opts.String("D2_LAYOUT", "layout", "l", "dagre", `the layout engine used`)
 	themeFlag, err := ms.Opts.Int64("D2_THEME", "theme", "t", 0, "the diagram theme ID")
 	if err != nil {
@@ -150,6 +154,9 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 	if *debugFlag {
 		ms.Env.Setenv("DEBUG", "1")
 	}
+	if *imgCacheFlag {
+		ms.Env.Setenv("IMG_CACHE", "1")
+	}
 	if *browserFlag != "" {
 		ms.Env.Setenv("BROWSER", *browserFlag)
 	}

lib/imgbundler/imgbundler.go

@@ -23,6 +23,8 @@ import (
 	"oss.terrastruct.com/util-go/xmain"
 )
 
+var imgCache sync.Map
+
 const maxImageSize int64 = 1 << 25 // 33_554_432
 
 var imageRegex = regexp.MustCompile(`<image href="([^"]+)"`)
@@ -55,12 +57,17 @@ func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_
 	return runWorkers(ctx, ms, svg, imgs, isRemote)
 }
 
-// filterImageElements finds all image elements in imgs that are eligible
-// for bundling in the current context.
+// filterImageElements finds all unique image elements in imgs that are
+// eligible for bundling in the current context.
 func filterImageElements(imgs [][][]byte, isRemote bool) [][][]byte {
+	unq := make(map[string]struct{})
 	imgs2 := imgs[:0]
 	for _, img := range imgs {
 		href := string(img[1])
+		if _, ok := unq[href]; ok {
+			continue
+		}
+		unq[href] = struct{}{}
 
 		// Skip already bundled images.
 		if strings.HasPrefix(href, "data:") {
@@ -104,7 +111,7 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
 					<-sema
 				}()
 
-				bundledImage, err := worker(ctx, img[1], isRemote)
+				bundledImage, err := worker(ctx, ms, img[1], isRemote)
 				if err != nil {
 					ms.Log.Error.Printf("failed to bundle %s: %v", img[1], err)
 					errhrefsMu.Lock()
@@ -138,18 +145,25 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
 				}
 				return svg, nil
 			}
-			svg = bytes.Replace(svg, repl.from, repl.to, 1)
+			svg = bytes.Replace(svg, repl.from, repl.to, -1)
 		}
 	}
 }
 
-func worker(ctx context.Context, href []byte, isRemote bool) ([]byte, error) {
+func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([]byte, error) {
+	if ms.Env.Getenv("IMG_CACHE") == "1" {
+		if hit, ok := imgCache.Load(string(href)); ok {
+			return hit.([]byte), nil
+		}
+	}
 	var buf []byte
 	var mimeType string
 	var err error
 	if isRemote {
+		ms.Log.Debug.Printf("fetching %s remotely", string(href))
 		buf, mimeType, err = httpGet(ctx, html.UnescapeString(string(href)))
 	} else {
+		ms.Log.Debug.Printf("reading %s from disk", string(href))
 		buf, err = os.ReadFile(html.UnescapeString(string(href)))
 	}
 	if err != nil {
@@ -161,7 +175,12 @@ func worker(ctx context.Context, href []byte, isRemote bool) ([]byte, error) {
 	}
 	mimeType = strings.Replace(mimeType, "text/xml", "image/svg+xml", 1)
 	b64 := base64.StdEncoding.EncodeToString(buf)
-	return []byte(fmt.Sprintf(`<image href="data:%s;base64,%s"`, mimeType, b64)), nil
+
+	out := []byte(fmt.Sprintf(`<image href="data:%s;base64,%s"`, mimeType, b64))
+	if ms.Env.Getenv("IMG_CACHE") == "1" {
+		imgCache.Store(string(href), out)
+	}
+	return out, nil
 }
 
 var httpClient = &http.Client{}

lib/imgbundler/imgbundler_test.go

@@ -10,8 +10,11 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 
+	tassert "github.com/stretchr/testify/assert"
+
 	"oss.terrastruct.com/util-go/cmdlog"
 	"oss.terrastruct.com/util-go/xos"
 
@@ -49,6 +52,7 @@ func TestRegex(t *testing.T) {
 }
 
 func TestInlineRemote(t *testing.T) {
+	imgCache = sync.Map{}
 	ctx := context.Background()
 	svgURL := "https://icons.terrastruct.com/essentials/004-picture.svg"
 	pngURL := "https://cdn4.iconfinder.com/data/icons/smart-phones-technologies/512/android-phone.png"
@@ -119,6 +123,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal("no png image inserted")
 	}
 
+	imgCache = sync.Map{}
 	// Test almost too large response
 	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
@@ -133,6 +138,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal(err)
 	}
 
+	imgCache = sync.Map{}
 	// Test too large response
 	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
@@ -147,6 +153,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal("expected error")
 	}
 
+	imgCache = sync.Map{}
 	// Test error response
 	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
@@ -160,6 +167,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 }
 
 func TestInlineLocal(t *testing.T) {
+	imgCache = sync.Map{}
 	ctx := context.Background()
 	svgURL, err := filepath.Abs("./test_svg.svg")
 	if err != nil {
@@ -221,3 +229,149 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal("no png image inserted")
 	}
 }
+
+// TestDuplicateURL ensures that we don't fetch the same image twice
+func TestDuplicateURL(t *testing.T) {
+	imgCache = sync.Map{}
+	ctx := context.Background()
+	url1 := "https://icons.terrastruct.com/essentials/004-picture.svg"
+	url2 := "https://icons.terrastruct.com/essentials/004-picture.svg"
+
+	sampleSVG := fmt.Sprintf(`<?xml version="1.0" encoding="utf-8"?>
+<svg
+id="d2-svg"
+style="background: white;"
+xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"
+width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
+<![CDATA[
+.shape {
+  shape-rendering: geometricPrecision;
+  stroke-linejoin: round;
+}
+.connection {
+  stroke-linecap: round;
+  stroke-linejoin: round;
+}
+
+]]>
+</style><g id="a"><g class="shape" ><image href="%s" x="0" y="0" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="-15.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">a</text></g><g id="b"><g class="shape" ><image href="%s" x="0" y="228" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="213.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">b</text></g><g id="(a -&gt; b)[0]"><marker id="mk-3990223579" markerWidth="10.000000" markerHeight="12.000000" refX="7.000000" refY="6.000000" viewBox="0.000000 0.000000 10.000000 12.000000" orient="auto" markerUnits="userSpaceOnUse"> <polygon class="connection" fill="#0D32B2" stroke-width="2" points="0.000000,0.000000 10.000000,6.000000 0.000000,12.000000" /> </marker><path d="M 64.000000 130.000000 C 64.000000 168.000000 64.000000 188.000000 64.000000 224.000000" class="connection" style="fill:none;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" marker-end="url(#mk-3990223579)" /></g><style type="text/css"><![CDATA[
+.text-bold {
+	font-family: "font-bold";
+}
+@font-face {
+	font-family: font-bold;
+	src: url("REMOVED");
+}]]></style></svg>
+`, url1, url2)
+
+	ms := &xmain.State{
+		Name: "test",
+
+		Stdin:  os.Stdin,
+		Stdout: os.Stdout,
+		Stderr: os.Stderr,
+
+		Env: xos.NewEnv(os.Environ()),
+	}
+	ms.Log = cmdlog.NewTB(ms.Env, t)
+
+	count := 0
+
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+		count++
+		respRecorder := httptest.NewRecorder()
+		respRecorder.WriteString(`<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\r\n<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->\r\n<svg version=\"1.1\" id=\"Capa_1\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" x=\"0px\" y=\"0px\"\r\n\t viewBox=\"0 0 58 58\" style=\"enable-background:new 0 0 58 58;\" xml:space=\"preserve\">\r\n<rect x=\"1\" y=\"7\" style=\"fill:#C3E1ED;stroke:#E7ECED;stroke-width:2;stroke-miterlimit:10;\" width=\"56\" height=\"44\"/>\r\n<circle style=\"fill:#ED8A19;\" cx=\"16\" cy=\"17.569\" r=\"6.569\"/>\r\n<polygon style=\"fill:#1A9172;\" points=\"56,36.111 55,35 43,24 32.5,35.5 37.983,40.983 42,45 56,45 \"/>\r\n<polygon style=\"fill:#1A9172;\" points=\"2,49 26,49 21.983,44.983 11.017,34.017 2,41.956 \"/>\r\n<rect x=\"2\" y=\"45\" style=\"fill:#6B5B4B;\" width=\"54\" height=\"5\"/>\r\n<polygon style=\"fill:#25AE88;\" points=\"37.983,40.983 27.017,30.017 10,45 42,45 \"/>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n</svg>`)
+		respRecorder.WriteHeader(200)
+		return respRecorder.Result()
+	})
+
+	out, err := BundleRemote(ctx, ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	tassert.Equal(t, 1, count)
+	if strings.Contains(string(out), url1) {
+		t.Fatal("links still exist")
+	}
+	tassert.Equal(t, 2, strings.Count(string(out), "image/svg+xml"))
+}
+
+func TestImgCache(t *testing.T) {
+	imgCache = sync.Map{}
+	ctx := context.Background()
+	url1 := "https://icons.terrastruct.com/essentials/004-picture.svg"
+	url2 := "https://icons.terrastruct.com/essentials/004-picture.svg"
+
+	sampleSVG := fmt.Sprintf(`<?xml version="1.0" encoding="utf-8"?>
+<svg
+id="d2-svg"
+style="background: white;"
+xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"
+width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
+<![CDATA[
+.shape {
+  shape-rendering: geometricPrecision;
+  stroke-linejoin: round;
+}
+.connection {
+  stroke-linecap: round;
+  stroke-linejoin: round;
+}
+
+]]>
+</style><g id="a"><g class="shape" ><image href="%s" x="0" y="0" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="-15.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">a</text></g><g id="b"><g class="shape" ><image href="%s" x="0" y="228" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="213.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">b</text></g><g id="(a -&gt; b)[0]"><marker id="mk-3990223579" markerWidth="10.000000" markerHeight="12.000000" refX="7.000000" refY="6.000000" viewBox="0.000000 0.000000 10.000000 12.000000" orient="auto" markerUnits="userSpaceOnUse"> <polygon class="connection" fill="#0D32B2" stroke-width="2" points="0.000000,0.000000 10.000000,6.000000 0.000000,12.000000" /> </marker><path d="M 64.000000 130.000000 C 64.000000 168.000000 64.000000 188.000000 64.000000 224.000000" class="connection" style="fill:none;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" marker-end="url(#mk-3990223579)" /></g><style type="text/css"><![CDATA[
+.text-bold {
+	font-family: "font-bold";
+}
+@font-face {
+	font-family: font-bold;
+	src: url("REMOVED");
+}]]></style></svg>
+`, url1, url2)
+
+	ms := &xmain.State{
+		Name: "test",
+
+		Stdin:  os.Stdin,
+		Stdout: os.Stdout,
+		Stderr: os.Stderr,
+
+		Env: xos.NewEnv(os.Environ()),
+	}
+	ms.Log = cmdlog.NewTB(ms.Env, t)
+
+	count := 0
+
+	httpClient.Transport = roundTripFunc(func(req *http.Request) *http.Response {
+		count++
+		respRecorder := httptest.NewRecorder()
+		respRecorder.WriteString(`<?xml version=\"1.0\" encoding=\"iso-8859-1\"?>\r\n<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->\r\n<svg version=\"1.1\" id=\"Capa_1\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" x=\"0px\" y=\"0px\"\r\n\t viewBox=\"0 0 58 58\" style=\"enable-background:new 0 0 58 58;\" xml:space=\"preserve\">\r\n<rect x=\"1\" y=\"7\" style=\"fill:#C3E1ED;stroke:#E7ECED;stroke-width:2;stroke-miterlimit:10;\" width=\"56\" height=\"44\"/>\r\n<circle style=\"fill:#ED8A19;\" cx=\"16\" cy=\"17.569\" r=\"6.569\"/>\r\n<polygon style=\"fill:#1A9172;\" points=\"56,36.111 55,35 43,24 32.5,35.5 37.983,40.983 42,45 56,45 \"/>\r\n<polygon style=\"fill:#1A9172;\" points=\"2,49 26,49 21.983,44.983 11.017,34.017 2,41.956 \"/>\r\n<rect x=\"2\" y=\"45\" style=\"fill:#6B5B4B;\" width=\"54\" height=\"5\"/>\r\n<polygon style=\"fill:#25AE88;\" points=\"37.983,40.983 27.017,30.017 10,45 42,45 \"/>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n<g>\r\n</g>\r\n</svg>`)
+		respRecorder.WriteHeader(200)
+		return respRecorder.Result()
+	})
+
+	// Using a cache, imgs are not refetched on multiple runs
+	ms.Env.Setenv("IMG_CACHE", "1")
+	_, err := BundleRemote(ctx, ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	tassert.Equal(t, 1, count)
+
+	// With cache disabled, it refetches
+	ms.Env.Setenv("IMG_CACHE", "0")
+	count = 0
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = BundleRemote(ctx, ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	tassert.Equal(t, 2, count)
+}