diff --git a/pkg/octicons/octicons.go b/pkg/octicons/octicons.go index 5954a8c223..7f93a2037a 100644 --- a/pkg/octicons/octicons.go +++ b/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 { diff --git a/pkg/octicons/octicons_benchmark_test.go b/pkg/octicons/octicons_benchmark_test.go new file mode 100644 index 0000000000..6c6cdd065d --- /dev/null +++ b/pkg/octicons/octicons_benchmark_test.go @@ -0,0 +1,74 @@ +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) + } + }) +} + +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) + } + } + + batch := make([][]mcp.Icon, len(inventory)) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + if !warm { + b.StopTimer() + dataURIs.mu.Lock() + dataURIs.values = nil + dataURIs.mu.Unlock() + b.StartTimer() + } + for index, icon := range inventory { + batch[index] = Icons(icon) + } + } + b.StopTimer() + benchmarkIconsSink = batch + runtime.KeepAlive(batch) + }) + } +} diff --git a/pkg/octicons/octicons_test.go b/pkg/octicons/octicons_test.go index 078eb744f2..5d08dec228 100644 --- a/pkg/octicons/octicons_test.go +++ b/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)