use context in indexer go routine instead of context.TODO (#1491)

* fixed bug of unprotected concurrent writes to a map in prefix indexer

Signed-off-by: Nir Rozenbaum <[email protected]>

* revert locking changes in indexer

Signed-off-by: Nir Rozenbaum <[email protected]>

* replace context.TODO with a context

Signed-off-by: Nir Rozenbaum <[email protected]>

---------

Signed-off-by: Nir Rozenbaum <[email protected]>

Author: nirrozenbaum

pkg/epp/scheduling/framework/plugins/multi/prefix/indexer.go

@@ -38,15 +38,15 @@ type indexer struct {
 }
 
 // newIndexer initializes an indexer with size limits and starts cache size reporting.
-func newIndexer(maxLRUSize int) *indexer {
-	ix := &indexer{
+func newIndexer(ctx context.Context, maxLRUSize int) *indexer {
+	indexer := &indexer{
 		hashToPods: make(map[BlockHash]podSet),
 		podToLRU:   make(map[ServerID]*lru.Cache[BlockHash, struct{}]),
 		maxLRUSize: maxLRUSize,
 	}
 
-	go ix.ReportLRUSize(time.Second)
-	return ix
+	go indexer.reportLRUSize(ctx, time.Second)
+	return indexer
 }
 
 // Add adds a list of prefix hashes to the cache, tied to the server.
@@ -111,8 +111,8 @@ func (i *indexer) makeEvictionFn(pod ServerID) func(BlockHash, struct{}) {
 	}
 }
 
-// ReportLRUSize starts a goroutine that periodically reports the LRU cache size metric.
-func (i *indexer) ReportLRUSize(interval time.Duration) {
+// reportLRUSize starts a goroutine that periodically reports the LRU cache size metric.
+func (i *indexer) reportLRUSize(ctx context.Context, interval time.Duration) {
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 	for range ticker.C {
@@ -137,7 +137,7 @@ func (i *indexer) ReportLRUSize(interval time.Duration) {
 		}
 
 		metrics.RecordPrefixCacheSize(int64(totalEntries))
-		log.FromContext(context.TODO()).V(logutil.TRACE).Info("Prefix cache state",
+		log.FromContext(ctx).V(logutil.TRACE).Info("Prefix cache state",
 			"total entries", totalEntries,
 			"# pods", numPods,
 			"avg entries per pod", avg,

pkg/epp/scheduling/framework/plugins/multi/prefix/indexer_test.go

@@ -16,13 +16,14 @@ limitations under the License.
 package prefix
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestIndexer_AddAndGet(t *testing.T) {
-	i := newIndexer(2)
+	i := newIndexer(context.Background(), 2)
 
 	hash1 := BlockHash(1)
 	server := ServerID{Namespace: "default", Name: "server1"}

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -148,7 +148,7 @@ func New(ctx context.Context, config Config) *Plugin {
 	capacity := config.LRUCapacityPerServer
 	if capacity <= 0 {
 		capacity = DefaultLRUCapacityPerServer
-		log.FromContext(context.TODO()).V(logutil.DEFAULT).Info(
+		log.FromContext(ctx).V(logutil.DEFAULT).Info(
 			"LRUCapacityPerServer is not positive, using default value",
 			"defaultCapacity", DefaultLRUCapacityPerServer,
 		)
@@ -158,7 +158,7 @@ func New(ctx context.Context, config Config) *Plugin {
 		typedName:   plugins.TypedName{Type: PrefixCachePluginType, Name: PrefixCachePluginType},
 		config:      config,
 		pluginState: plugins.NewPluginState(ctx),
-		indexer:     newIndexer(capacity),
+		indexer:     newIndexer(ctx, capacity),
 	}
 }