perf(octicons): cache embedded data URIs on demand

Author: dexhunter

pkg/octicons/octicons.go

@@ -8,6 +8,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"strings"
+	"sync"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
@@ -18,6 +19,47 @@ var iconsFS embed.FS
 //go:embed required_icons.txt
 var requiredIconsTxt string
 
+type dataURIKey struct {
+	name  string
+	theme Theme
+}
+
+type dataURICache struct {
+	mu     sync.RWMutex
+	values map[dataURIKey]string
+}
+
+// This covers the current embedded variants without encoding them at startup.
+const initialDataURICacheCapacity = 64
+
+func (c *dataURICache) load(name string, theme Theme) string {
+	key := dataURIKey{name: name, theme: theme}
+	c.mu.RLock()
+	cached, ok := c.values[key]
+	c.mu.RUnlock()
+	if ok {
+		return cached
+	}
+
+	dataURI := readDataURI(name, theme)
+	if dataURI == "" {
+		return ""
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if cached, ok := c.values[key]; ok {
+		return cached
+	}
+	if c.values == nil {
+		c.values = make(map[dataURIKey]string, initialDataURICacheCapacity)
+	}
+	c.values[key] = dataURI
+	return dataURI
+}
+
+var dataURIs dataURICache
+
 // RequiredIcons returns the list of icon names from required_icons.txt.
 // This is the single source of truth for which icons should be embedded.
 func RequiredIcons() []string {
@@ -50,6 +92,10 @@ const (
 // - ThemeDark: light icons for dark backgrounds
 // If the icon is not found in the embedded filesystem, it returns an empty string.
 func DataURI(name string, theme Theme) string {
+	return dataURIs.load(name, theme)
+}
+
+func readDataURI(name string, theme Theme) string {
 	filename := fmt.Sprintf("icons/%s-%s.png", name, theme)
 	data, err := iconsFS.ReadFile(filename)
 	if err != nil {

pkg/octicons/octicons_benchmark_test.go

@@ -0,0 +1,73 @@
+package octicons
+
+import (
+	"runtime"
+	"testing"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+var benchmarkDataURISink string
+var benchmarkIconsSink [][]mcp.Icon
+
+func BenchmarkDataURICache(b *testing.B) {
+	b.Run("cold", func(b *testing.B) {
+		b.ReportAllocs()
+		for b.Loop() {
+			var cache dataURICache
+			benchmarkDataURISink = cache.load("repo", ThemeLight)
+		}
+	})
+
+	b.Run("warm", func(b *testing.B) {
+		var cache dataURICache
+		benchmarkDataURISink = cache.load("repo", ThemeLight)
+		b.ReportAllocs()
+		b.ResetTimer()
+		for b.Loop() {
+			benchmarkDataURISink = cache.load("repo", ThemeLight)
+		}
+	})
+}
+
+// Run each ColdStart inventory separately with -benchtime=1x in a fresh
+// `go test` process so package-level cache state cannot cross-contaminate it.
+func BenchmarkIconsRegistrationColdStart(b *testing.B) {
+	benchmarkIconsRegistration(b, false)
+}
+
+func BenchmarkIconsRegistrationWarm(b *testing.B) {
+	benchmarkIconsRegistration(b, true)
+}
+
+func benchmarkIconsRegistration(b *testing.B, warm bool) {
+	inventories := map[string][]string{
+		"narrow":  {"repo"},
+		"default": RequiredIcons(),
+	}
+	for name, inventory := range inventories {
+		b.Run(name, func(b *testing.B) {
+			if warm {
+				for _, icon := range inventory {
+					_ = Icons(icon)
+				}
+			} else {
+				dataURIs.mu.Lock()
+				dataURIs.values = nil
+				dataURIs.mu.Unlock()
+			}
+
+			batch := make([][]mcp.Icon, len(inventory))
+			b.ReportAllocs()
+			b.ResetTimer()
+			for b.Loop() {
+				for index, icon := range inventory {
+					batch[index] = Icons(icon)
+				}
+			}
+			b.StopTimer()
+			benchmarkIconsSink = batch
+			runtime.KeepAlive(batch)
+		})
+	}
+}

pkg/octicons/octicons_test.go

@@ -1,7 +1,9 @@
 package octicons
 
 import (
+	"io/fs"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -53,6 +55,65 @@ func TestDataURI(t *testing.T) {
 	}
 }
 
+func TestDataURIForEveryEmbeddedIcon(t *testing.T) {
+	paths, err := fs.Glob(iconsFS, "icons/*.png")
+	assert.NoError(t, err)
+	assert.NotEmpty(t, paths)
+
+	for _, path := range paths {
+		filename := strings.TrimSuffix(strings.TrimPrefix(path, "icons/"), ".png")
+		separator := strings.LastIndexByte(filename, '-')
+		if separator <= 0 {
+			t.Errorf("cannot parse embedded icon path %q", path)
+			continue
+		}
+		name := filename[:separator]
+		theme := Theme(filename[separator+1:])
+		t.Run(filename, func(t *testing.T) {
+			assert.True(t, strings.HasPrefix(DataURI(name, theme), "data:image/png;base64,"))
+		})
+	}
+}
+
+func TestDataURICacheOnlyStoresSuccessfulReads(t *testing.T) {
+	var cache dataURICache
+	missingKey := dataURIKey{name: "nonexistent-icon", theme: ThemeLight}
+
+	assert.Empty(t, cache.load(missingKey.name, missingKey.theme))
+	assert.Nil(t, cache.values, "missing icons must not initialize the cache")
+	_, found := cache.values[missingKey]
+	assert.False(t, found, "missing icons must not be cached")
+
+	validKey := dataURIKey{name: "repo", theme: ThemeLight}
+	assert.NotEmpty(t, cache.load(validKey.name, validKey.theme))
+	_, found = cache.values[validKey]
+	assert.True(t, found, "successful reads should be cached")
+}
+
+func TestDataURICacheConcurrentFirstUse(t *testing.T) {
+	var cache dataURICache
+	want := readDataURI("repo", ThemeLight)
+	assert.NotEmpty(t, want)
+
+	const workers = 64
+	start := make(chan struct{})
+	results := make(chan string, workers)
+	var wg sync.WaitGroup
+	for range workers {
+		wg.Go(func() {
+			<-start
+			results <- cache.load("repo", ThemeLight)
+		})
+	}
+
+	close(start)
+	wg.Wait()
+	close(results)
+	for result := range results {
+		assert.Equal(t, want, result)
+	}
+}
+
 func TestIcons(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -99,6 +160,23 @@ func TestIcons(t *testing.T) {
 	}
 }
 
+func TestIconsReturnsFreshSlice(t *testing.T) {
+	lightSource := DataURI("repo", ThemeLight)
+	darkSource := DataURI("repo", ThemeDark)
+
+	icons := Icons("repo")
+	icons[0] = mcp.Icon{}
+	icons[1].Source = "mutated"
+
+	fresh := Icons("repo")
+	assert.Equal(t, lightSource, fresh[0].Source)
+	assert.Equal(t, darkSource, fresh[1].Source)
+	assert.Equal(t, "image/png", fresh[0].MIMEType)
+	assert.Equal(t, "image/png", fresh[1].MIMEType)
+	assert.Equal(t, mcp.IconThemeLight, fresh[0].Theme)
+	assert.Equal(t, mcp.IconThemeDark, fresh[1].Theme)
+}
+
 func TestThemeConstants(t *testing.T) {
 	assert.Equal(t, Theme("light"), ThemeLight)
 	assert.Equal(t, Theme("dark"), ThemeDark)