Update tests to match usage

core/pkg/util/stringutil/lrubank.go

@@ -139,7 +139,7 @@ func (sb *lruStringBank) LoadOrStoreFunc(key string, f func() string) (string, b
 
 	// create the key and value using the func (the key could be deallocated later)
 	value := f()
-	sb.m[key] = &lruEntry{
+	sb.m[value] = &lruEntry{
 		value: value,
 		used:  time.Now().UTC(),
 	}

core/pkg/util/stringutil/lrubank_test.go

@@ -30,6 +30,7 @@ func TestBasicLruEvict(t *testing.T) {
 			t.Errorf("The 'bar' entry should've been replaced by 'test'")
 		}
 	}
+	lruBank.lock.Unlock()
 }
 
 // ---------------------------------------------------------------------------
@@ -42,7 +43,7 @@ func TestLoadOrStore_MissAndHit(t *testing.T) {
 	bank := NewLruStringBank(10, time.Minute).(*lruStringBank)
 	defer bank.Stop()
 
-	v, loaded := bank.LoadOrStore("k", "hello")
+	v, loaded := bank.LoadOrStore("hello", "hello")
 	if loaded {
 		t.Errorf("first LoadOrStore: expected loaded=false, got true")
 	}
@@ -50,7 +51,7 @@ func TestLoadOrStore_MissAndHit(t *testing.T) {
 		t.Errorf("first LoadOrStore: expected value %q, got %q", "hello", v)
 	}
 
-	v, loaded = bank.LoadOrStore("k", "world")
+	v, loaded = bank.LoadOrStore("hello", "world")
 	if !loaded {
 		t.Errorf("second LoadOrStore: expected loaded=true, got false")
 	}
@@ -104,7 +105,7 @@ func TestLoadOrStoreFunc_FactoryCalledOnMissOnly(t *testing.T) {
 
 	factory := func() string {
 		calls++
-		return "generated"
+		return "k"
 	}
 
 	bank.LoadOrStoreFunc("k", factory)
@@ -115,32 +116,6 @@ func TestLoadOrStoreFunc_FactoryCalledOnMissOnly(t *testing.T) {
 	}
 }
 
-// Regression test: LoadOrStoreFunc stores entries under the *value* key, not
-// the *key* argument.  Subsequent lookups by the original key will therefore
-// miss every time.  This test documents the behaviour so any fix is caught
-// immediately.
-func TestLoadOrStoreFunc_KeyBug(t *testing.T) {
-	bank := NewLruStringBank(10, time.Minute).(*lruStringBank)
-	defer bank.Stop()
-
-	// First call: miss → stores entry under value "v", not key "k".
-	v, loaded := bank.LoadOrStoreFunc("k", func() string { return "v" })
-	if loaded {
-		t.Errorf("first call: expected loaded=false")
-	}
-	if v != "v" {
-		t.Errorf("first call: expected value %q, got %q", "v", v)
-	}
-
-	// Second call with the same key: because the entry was stored under "v",
-	// looking up "k" misses again.  This is the bug: loaded should be true
-	// but the current implementation returns false.
-	_, loaded = bank.LoadOrStoreFunc("k", func() string { return "v" })
-	if !loaded {
-		t.Errorf("LoadOrStoreFunc second call returned loaded=false")
-	}
-}
-
 // ---------------------------------------------------------------------------
 // Capacity / eviction
 // ---------------------------------------------------------------------------
@@ -152,7 +127,7 @@ func TestEviction_BelowCapacityNoEviction(t *testing.T) {
 	defer bank.Stop()
 
 	for i := 0; i < capacity; i++ {
-		bank.LoadOrStore(fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i))
+		bank.LoadOrStore(fmt.Sprintf("v%d", i), fmt.Sprintf("v%d", i))
 	}
 
 	// Wait several eviction cycles.
@@ -173,7 +148,7 @@ func TestEviction_ExceedCapacityTrimsToCapacity(t *testing.T) {
 	defer bank.Stop()
 
 	for i := 0; i < capacity+3; i++ {
-		bank.LoadOrStore(fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i))
+		bank.LoadOrStore(fmt.Sprintf("v%d", i), fmt.Sprintf("v%d", i))
 		time.Sleep(20 * time.Millisecond) // ensure distinct timestamps
 	}
 
@@ -225,7 +200,7 @@ func TestClear_EmptiesMap(t *testing.T) {
 	defer bank.Stop()
 
 	for i := 0; i < 5; i++ {
-		bank.LoadOrStore(fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i))
+		bank.LoadOrStore(fmt.Sprintf("v%d", i), fmt.Sprintf("v%d", i))
 	}
 
 	bank.Clear()
@@ -333,20 +308,33 @@ func TestConcurrentLoadOrStore(t *testing.T) {
 	const opsEach = 100
 
 	var wg sync.WaitGroup
-	wg.Add(goroutines)
-
-	for g := 0; g < goroutines; g++ {
-		g := g
-		go func() {
-			defer wg.Done()
+	for i := 0; i < goroutines; i++ {
+		g := i
+		wg.Go(func() {
 			for i := 0; i < opsEach; i++ {
 				key := fmt.Sprintf("k%d", (g*opsEach+i)%30)
 				bank.LoadOrStore(key, key)
 			}
+		})
+	}
+
+	waiter := func() chan struct{} {
+		st := make(chan struct{})
+
+		go func() {
+			wg.Wait()
+			close(st)
 		}()
+
+		return st
 	}
 
-	wg.Wait()
+	select {
+	case <-waiter():
+		t.Logf("Completed Successfully\n")
+	case <-time.After(10 * time.Second):
+		t.Logf("Timed out\n")
+	}
 }
 
 // Concurrent calls interleaved with eviction cycles must not deadlock or race.
@@ -358,13 +346,12 @@ func TestConcurrentLoadOrStoreWithEviction(t *testing.T) {
 	const duration = 300 * time.Millisecond
 
 	var wg sync.WaitGroup
-	stop := time.After(duration)
 
 	for i := 0; i < goroutines; i++ {
 		g := i
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
+		stop := time.After(duration)
+
+		wg.Go(func() {
 			for {
 				select {
 				case <-stop:
@@ -374,10 +361,26 @@ func TestConcurrentLoadOrStoreWithEviction(t *testing.T) {
 					bank.LoadOrStore(key, key)
 				}
 			}
+		})
+	}
+
+	waiter := func() chan struct{} {
+		st := make(chan struct{})
+
+		go func() {
+			wg.Wait()
+			close(st)
 		}()
+
+		return st
 	}
 
-	wg.Wait()
+	select {
+	case <-waiter():
+		t.Logf("Completed Successfully\n")
+	case <-time.After(10 * time.Second):
+		t.Logf("Timed out\n")
+	}
 }
 
 // Concurrent Clear calls alongside reads/writes must not panic.

core/pkg/util/stringutil/mapbank.go

@@ -36,7 +36,7 @@ func (sb *stringBank) LoadOrStoreFunc(key string, f func() string) (string, bool
 
 	// create the key and value using the func (the key could be deallocated later)
 	value := f()
-	sb.m[key] = value
+	sb.m[value] = value
 	sb.lock.Unlock()
 	return value, false
 }

core/pkg/util/stringutil/stringutil_test.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 	"testing"
 	"time"
+	"unsafe"
 
 	"github.com/opencost/opencost/core/pkg/util/stringutil"
 )
@@ -47,7 +48,7 @@ func copyString(s string) string {
 	return string([]byte(s))
 }
 
-func generateBenchData(totalStrings, totalUnique int) []string {
+func generateBenchData(totalStrings, totalUnique int) [][]byte {
 	randStrings := make([]string, 0, totalStrings)
 	r := rand.New(rand.NewSource(27644437))
 
@@ -70,7 +71,11 @@ func generateBenchData(totalStrings, totalUnique int) []string {
 	// shuffle the list of strings
 	r.Shuffle(totalStrings, func(i, j int) { randStrings[i], randStrings[j] = randStrings[j], randStrings[i] })
 
-	return randStrings
+	stringBytes := make([][]byte, 0, totalStrings)
+	for _, str := range randStrings {
+		stringBytes = append(stringBytes, []byte(str))
+	}
+	return stringBytes
 }
 
 func benchmarkStringBank(b *testing.B, bt bankTest, totalStrings, totalUnique int, useBankFunc bool) {
@@ -81,10 +86,16 @@ func benchmarkStringBank(b *testing.B, bt bankTest, totalStrings, totalUnique in
 		for i := 0; i < b.N; i++ {
 			b.StartTimer()
 			for bb := 0; bb < totalStrings; bb++ {
+				bytes := randStrings[bb]
+
 				if useBankFunc {
-					bt.BankFunc(randStrings[bb], func() string { return randStrings[bb] })
+					str := unsafe.String(unsafe.SliceData(bytes), len(bytes))
+
+					bt.BankFunc(str, func() string {
+						return string(bytes)
+					})
 				} else {
-					bt.Bank(randStrings[bb])
+					bt.Bank(string(bytes))
 				}
 			}
 			b.StopTimer()