fix: use atomic increment for image ID allocation to prevent race condition (#748)

Daniel Purnomo created 2 months ago

Change summary

view/html.go      |  5 ++---
view/html_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
2 files changed, 45 insertions(+), 3 deletions(-)

Detailed changes

view/html.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"charm.land/lipgloss/v2"
@@ -282,9 +283,7 @@ var nextImageID uint32 = 1000
 
 // allocImageID returns a unique Kitty image ID.
 func allocImageID() uint32 {
-	id := nextImageID
-	nextImageID++
-	return id
+	return atomic.AddUint32(&nextImageID, 1)
 }
 
 func fetchRemoteBase64(url string) string {

view/html_test.go

@@ -5,6 +5,8 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"charm.land/lipgloss/v2"
@@ -761,3 +763,44 @@ func TestRemoteImageCache_EvictsOldestWhenFull(t *testing.T) {
 		}
 	}
 }
+
+func TestAllocImageID_NoRace(t *testing.T) {
+	// Reset the counter so IDs start from a known value.
+	atomic.StoreUint32(&nextImageID, 1000)
+
+	const goroutines = 100
+	const idsPerGoroutine = 100
+
+	results := make(chan uint32, goroutines*idsPerGoroutine)
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	for range goroutines {
+		go func() {
+			defer wg.Done()
+			for range idsPerGoroutine {
+				results <- allocImageID()
+			}
+		}()
+	}
+
+	// Close channel once all writers are done.
+	go func() {
+		wg.Wait()
+		close(results)
+	}()
+
+	// Collect all IDs and verify uniqueness.
+	seen := make(map[uint32]bool, goroutines*idsPerGoroutine)
+	for id := range results {
+		if seen[id] {
+			t.Fatalf("duplicate image ID allocated: %d (race condition detected)", id)
+		}
+		seen[id] = true
+	}
+
+	expected := uint32(goroutines * idsPerGoroutine)
+	if uint32(len(seen)) != expected {
+		t.Errorf("expected %d unique IDs, got %d", expected, len(seen))
+	}
+}