fix: clear KV block index entries on endpoint removal (#1738)

* fix: clear KV block index entries on endpoint removal

On EventDelete, call Index.Clear with the pod's IP address to
remove stale entries that reference blocks no longer cached on
the removed pod.

Fixes #1736

Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>

* Update pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/extractor.go

Co-authored-by: Bongwoo Bak <bongwoobak@gmail.com>
Signed-off-by: Rishabh Saini <rishabhsaini01@gmail.com>

* fix(test): match Clear assertion to IP:Port PodIdentifier format

Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>

* Update pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/extractor.go

Co-authored-by: Maroon Ayoub <Maroonay@gmail.com>
Signed-off-by: Rishabh Saini <rishabhsaini01@gmail.com>

* Update pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/extractor.go

Co-authored-by: Maroon Ayoub <Maroonay@gmail.com>
Signed-off-by: Rishabh Saini <rishabhsaini01@gmail.com>

---------

Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
Signed-off-by: Rishabh Saini <rishabhsaini01@gmail.com>
Co-authored-by: Bongwoo Bak <bongwoobak@gmail.com>
Co-authored-by: Maroon Ayoub <Maroonay@gmail.com>

Author: 3 people

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/extractor.go

@@ -52,6 +52,12 @@ func (p *Producer) Extract(ctx context.Context, event fwkdl.EndpointEvent) error
 		logger.V(logging.DEBUG).Info("Adding subscriber", "endpoint", endpointKey)
 	case fwkdl.EventDelete:
 		p.subscribersManager.RemoveSubscriber(ctx, endpointKey)
+		if meta.Address != "" {
+			if err := p.kvCacheIndexer.KVBlockIndex().Clear(ctx, fmt.Sprintf("%s:%s", meta.Address, meta.Port)); err != nil {
+				logger.Error(err, "Failed to clear index entries for removed endpoint",
+					"endpoint", endpointKey, "address", meta.Address, "port", meta.Port)
+			}
+		}
 		logger.V(logging.DEBUG).Info("Removed KV-events subscriber", "endpoint", endpointKey)
 	}
 	return nil

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/extractor_test.go

@@ -49,6 +49,7 @@ func newExtractorProducer(discoverPods bool) *Producer {
 		typedName:          plugin.TypedName{Type: PluginType, Name: PluginType},
 		subscribersManager: kvevents.NewSubscriberManager(kvevents.NewPool(cfg, nil, nil, nil)),
 		kvEventsConfig:     cfg,
+		kvCacheIndexer:     &fakeKVCacheIndexer{index: &fakeKVBlockIndex{}},
 		subscriberCtx:      context.Background(),
 	}
 }
@@ -216,6 +217,51 @@ func TestProducer_ExtractEndpoint_SingleRankUsesBaseSocketPort(t *testing.T) {
 		"single-rank pod (RankIndex=0) must dial the base SocketPort")
 }
 
+// EventDelete clears index entries for the removed pod's address.
+func TestProducer_ExtractEndpoint_DeleteClearsIndex(t *testing.T) {
+	ctx := discardCtx(t)
+
+	var clearedPod string
+	fakeIndex := &fakeKVBlockIndex{
+		clearFn: func(_ context.Context, podIdentifier string) error {
+			clearedPod = podIdentifier
+			return nil
+		},
+	}
+	fakeIndexer := &fakeKVCacheIndexer{index: fakeIndex}
+
+	cfg := kvevents.DefaultConfig()
+	cfg.DiscoverPods = true
+	cfg.PodDiscoveryConfig = kvevents.DefaultPodReconcilerConfig()
+	cfg.PodDiscoveryConfig.SocketPort = 5557
+
+	p := &Producer{
+		typedName:          plugin.TypedName{Type: PluginType, Name: PluginType},
+		subscribersManager: kvevents.NewSubscriberManager(kvevents.NewPool(cfg, nil, nil, nil)),
+		kvEventsConfig:     cfg,
+		kvCacheIndexer:     fakeIndexer,
+		subscriberCtx:      context.Background(),
+	}
+	defer p.subscribersManager.Shutdown(ctx)
+
+	ep := newEndpoint("pod-clear", "10.0.0.99")
+
+	require.NoError(t, p.Extract(ctx, fwkdl.EndpointEvent{
+		Type:     fwkdl.EventAddOrUpdate,
+		Endpoint: ep,
+	}))
+
+	require.NoError(t, p.Extract(ctx, fwkdl.EndpointEvent{
+		Type:     fwkdl.EventDelete,
+		Endpoint: ep,
+	}))
+
+	assert.Equal(t, "10.0.0.99:8080", clearedPod, "index should be cleared using pod IP:Port matching PodIdentifier format")
+
+	ids, _ := p.subscribersManager.GetActiveSubscribers()
+	assert.Empty(t, ids)
+}
+
 // Delete by NamespacedName must work even when the event has no address.
 func TestProducer_ExtractEndpoint_DeleteWithMissingAddressRemovesExistingSubscriber(t *testing.T) {
 	ctx := discardCtx(t)

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go

@@ -52,8 +52,9 @@ func (f *fakeKVCacheIndexer) ComputeBlockKeysFromTokens(ctx context.Context, tok
 func (f *fakeKVCacheIndexer) KVBlockIndex() kvblock.Index { return f.index }
 
 type fakeKVBlockIndex struct {
-	lookup func(ctx context.Context, keys []kvblock.BlockHash, podSet sets.Set[string]) (map[kvblock.BlockHash][]kvblock.PodEntry, error)
-	addFn  func(ctx context.Context, prevKeys, keys []kvblock.BlockHash, entries []kvblock.PodEntry) error
+	lookup  func(ctx context.Context, keys []kvblock.BlockHash, podSet sets.Set[string]) (map[kvblock.BlockHash][]kvblock.PodEntry, error)
+	addFn   func(ctx context.Context, prevKeys, keys []kvblock.BlockHash, entries []kvblock.PodEntry) error
+	clearFn func(ctx context.Context, podIdentifier string) error
 }
 
 func (f *fakeKVBlockIndex) Lookup(ctx context.Context, keys []kvblock.BlockHash, podSet sets.Set[string]) (map[kvblock.BlockHash][]kvblock.PodEntry, error) {
@@ -78,7 +79,10 @@ func (f *fakeKVBlockIndex) GetRequestKey(_ context.Context, _ kvblock.BlockHash)
 	return kvblock.EmptyBlockHash, nil
 }
 
-func (f *fakeKVBlockIndex) Clear(_ context.Context, _ string) error {
+func (f *fakeKVBlockIndex) Clear(ctx context.Context, podIdentifier string) error {
+	if f.clearFn != nil {
+		return f.clearFn(ctx, podIdentifier)
+	}
 	return nil
 }