Update Index to retain UID across scrapes (#4065)

Signed-off-by: Sean Holcomb <seanholcomb@gmail.com>

modules/collector-source/pkg/scrape/clustercache.go

@@ -28,12 +28,20 @@ const unmountedPVsContainer = "unmounted-pvs"
 type ClusterCacheScraper struct {
 	clusterCache          clustercache.ClusterCache
 	externalLabelProvider external.LabelProvider
+	nodeIndex             *persistedIndex[string]
+	namespaceIndex        *persistedIndex[string]
+	pvcIndex              *persistedIndex[pvcKey]
+	pvIndex               *persistedIndex[string]
 }
 
 func newClusterCacheScraper(clusterCache clustercache.ClusterCache, externalLabelProvider external.LabelProvider) Scraper {
 	return &ClusterCacheScraper{
 		clusterCache:          clusterCache,
 		externalLabelProvider: externalLabelProvider,
+		nodeIndex:             newPersistedIndex[string]("node"),
+		namespaceIndex:        newPersistedIndex[string]("namespace"),
+		pvcIndex:              newPersistedIndex[pvcKey]("pvc"),
+		pvIndex:               newPersistedIndex[string]("pv"),
 	}
 }
 
@@ -54,11 +62,12 @@ func (ccs *ClusterCacheScraper) Scrape() []metric.Update {
 	resourceQuotas := ccs.clusterCache.GetAllResourceQuotas()
 
 	// create scrape indexes. While the pairs being mapped here don't have a 1 to 1 relationship in the general case,
-	// we are assuming that in the context of a single snapshot of the cluster they are 1 to 1.
-	nodeNameToUID := buildNodeIndex(nodes)
-	namespaceNameToUID := buildNamespaceIndex(namespaces)
-	pvcNameToUID := buildPVCIndex(pvcs)
-	pvNameToUID := buildPVIndex(pvs)
+	// we are assuming that in the context of a single snapshot of the cluster they are 1 to 1. Entries are retained
+	// across scrapes so that objects which outlive their referent in the cluster cache still resolve a UID.
+	nodeNameToUID := ccs.nodeIndex.update(buildNodeIndex(nodes))
+	namespaceNameToUID := ccs.namespaceIndex.update(buildNamespaceIndex(namespaces))
+	pvcNameToUID := ccs.pvcIndex.update(buildPVCIndex(pvcs))
+	pvNameToUID := ccs.pvIndex.update(buildPVIndex(pvs))
 
 	scrapeFuncs := []ScrapeFunc{
 		ccs.GetScrapeNodes(nodes),

modules/collector-source/pkg/scrape/clustercache_test.go

@@ -3092,3 +3092,46 @@ func Test_kubernetesScraper_scrapeCronJobs(t *testing.T) {
 		})
 	}
 }
+
+func TestClusterCacheScraper_Scrape_PodRetainsUIDsOfRemovedReferents(t *testing.T) {
+	cache := &clustercache.MockClusterCache{
+		Nodes:      []*clustercache.Node{{Name: "node-a", UID: "uid-node-a"}},
+		Namespaces: []*clustercache.Namespace{{Name: "ns-1", UID: "uid-ns-1"}},
+		Pods: []*clustercache.Pod{
+			{
+				Name:      "pod-a",
+				Namespace: "ns-1",
+				UID:       "uid-pod-a",
+				Spec:      clustercache.PodSpec{NodeName: "node-a"},
+			},
+		},
+	}
+	ccs := newClusterCacheScraper(cache, nil).(*ClusterCacheScraper)
+
+	podInfo := func(updates []metric.Update) map[string]string {
+		for _, u := range updates {
+			if u.Name == metric.PodInfo {
+				return u.AdditionalInfo
+			}
+		}
+		t.Fatalf("no %s update found", metric.PodInfo)
+		return nil
+	}
+
+	first := podInfo(ccs.Scrape())
+	if first[source.NodeUIDLabel] != "uid-node-a" || first[source.NamespaceUIDLabel] != "uid-ns-1" {
+		t.Fatalf("unexpected pod info on first scrape: %v", first)
+	}
+
+	// the node and namespace are removed from the cluster cache before the pod
+	cache.Nodes = nil
+	cache.Namespaces = nil
+
+	second := podInfo(ccs.Scrape())
+	if second[source.NodeUIDLabel] != "uid-node-a" {
+		t.Errorf("expected node UID to be retained, got %q", second[source.NodeUIDLabel])
+	}
+	if second[source.NamespaceUIDLabel] != "uid-ns-1" {
+		t.Errorf("expected namespace UID to be retained, got %q", second[source.NamespaceUIDLabel])
+	}
+}

modules/collector-source/pkg/scrape/index.go

@@ -1,10 +1,70 @@
 package scrape
 
 import (
+	"sync"
+	"time"
+
 	"github.com/opencost/opencost/core/pkg/clustercache"
+	"github.com/opencost/opencost/core/pkg/log"
 	"k8s.io/apimachinery/pkg/types"
 )
 
+// persistedIndexTTL is how long an index entry is retained after it was last seen in the cluster cache. It needs to
+// outlast the window where dependent objects (e.g. pods on a deleted node) remain in the cache after their referent
+// has been removed.
+const persistedIndexTTL = time.Hour
+
+type persistedIndexEntry struct {
+	uid      types.UID
+	lastSeen time.Time
+}
+
+// persistedIndex retains key to UID mappings across scrapes. Objects are not always removed from the cluster cache in
+// dependency order, so an object can outlive the object it references by name (e.g. a pod whose node has already been
+// deleted). Retaining recently seen entries prevents these lookups from resolving to an empty UID, which would
+// overwrite the previously scraped value. A nil *persistedIndex performs no retention.
+type persistedIndex[K comparable] struct {
+	name    string
+	lock    sync.Mutex
+	entries map[K]persistedIndexEntry
+}
+
+func newPersistedIndex[K comparable](name string) *persistedIndex[K] {
+	return &persistedIndex[K]{
+		name:    name,
+		entries: make(map[K]persistedIndexEntry),
+	}
+}
+
+// update records the entries of current, which always take precedence over retained entries, evicts entries which
+// have not been seen within the ttl, and returns a new index containing current along with any retained entries.
+func (pi *persistedIndex[K]) update(current map[K]types.UID) map[K]types.UID {
+	if pi == nil {
+		return current
+	}
+
+	pi.lock.Lock()
+	defer pi.lock.Unlock()
+
+	now := time.Now()
+	for key, uid := range current {
+		pi.entries[key] = persistedIndexEntry{uid: uid, lastSeen: now}
+	}
+
+	result := make(map[K]types.UID, len(pi.entries))
+	for key, entry := range pi.entries {
+		if now.Sub(entry.lastSeen) > persistedIndexTTL {
+			delete(pi.entries, key)
+			continue
+		}
+		if _, ok := current[key]; !ok {
+			log.Debugf("%s index: retaining UID '%s' for '%v' which is no longer in the cluster cache", pi.name, entry.uid, key)
+		}
+		result[key] = entry.uid
+	}
+	return result
+}
+
 // pvcKey is a composite key for a PersistentVolumeClaim (name + namespace).
 type pvcKey struct {
 	name      string

modules/collector-source/pkg/scrape/index_test.go

@@ -2,6 +2,7 @@ package scrape
 
 import (
 	"testing"
+	"time"
 
 	"github.com/opencost/opencost/core/pkg/clustercache"
 	"github.com/stretchr/testify/require"
@@ -73,3 +74,57 @@ func TestBuildPVIndex(t *testing.T) {
 	require.Equal(t, types.UID("uid-pv-b"), m["pv-b"])
 	require.Len(t, m, 2)
 }
+
+// backdate moves the last seen time of an entry into the past.
+func backdate[K comparable](pi *persistedIndex[K], key K, d time.Duration) {
+	entry := pi.entries[key]
+	entry.lastSeen = entry.lastSeen.Add(-d)
+	pi.entries[key] = entry
+}
+
+func TestPersistedIndex_Nil(t *testing.T) {
+	var pi *persistedIndex[string]
+	current := map[string]types.UID{"node-a": "uid-a"}
+	require.Equal(t, current, pi.update(current))
+}
+
+func TestPersistedIndex_RetainsMissingEntries(t *testing.T) {
+	pi := newPersistedIndex[string]("test")
+
+	m := pi.update(map[string]types.UID{"node-a": "uid-a", "node-b": "uid-b"})
+	require.Equal(t, map[string]types.UID{"node-a": "uid-a", "node-b": "uid-b"}, m)
+
+	backdate(pi, "node-b", persistedIndexTTL/2)
+	m = pi.update(map[string]types.UID{"node-a": "uid-a"})
+	require.Equal(t, map[string]types.UID{"node-a": "uid-a", "node-b": "uid-b"}, m)
+}
+
+func TestPersistedIndex_CurrentOverwritesRetained(t *testing.T) {
+	pi := newPersistedIndex[string]("test")
+
+	pi.update(map[string]types.UID{"node-a": "uid-a-old"})
+	m := pi.update(map[string]types.UID{"node-a": "uid-a-new"})
+	require.Equal(t, map[string]types.UID{"node-a": "uid-a-new"}, m)
+}
+
+func TestPersistedIndex_EvictsExpiredEntries(t *testing.T) {
+	pi := newPersistedIndex[string]("test")
+
+	pi.update(map[string]types.UID{"node-a": "uid-a", "node-b": "uid-b"})
+
+	backdate(pi, "node-b", 2*persistedIndexTTL)
+	m := pi.update(map[string]types.UID{"node-a": "uid-a"})
+	require.Equal(t, map[string]types.UID{"node-a": "uid-a"}, m)
+	require.Len(t, pi.entries, 1)
+}
+
+func TestPersistedIndex_ResultIsIndependent(t *testing.T) {
+	pi := newPersistedIndex[pvcKey]("test")
+
+	key := pvcKey{name: "pvc-a", namespace: "ns-1"}
+	m := pi.update(map[pvcKey]types.UID{key: "uid-pvc-a"})
+	m[key] = "modified"
+
+	m = pi.update(nil)
+	require.Equal(t, types.UID("uid-pvc-a"), m[key])
+}

modules/collector-source/pkg/scrape/statsummary.go

@@ -14,19 +14,23 @@ import (
 type StatSummaryScraper struct {
 	client       nodestats.StatSummaryClient
 	clusterCache clustercache.ClusterCache
+	nodeIndex    *persistedIndex[string]
+	pvcIndex     *persistedIndex[pvcKey]
 }
 
 func newStatSummaryScraper(client nodestats.StatSummaryClient, clusterCache clustercache.ClusterCache) Scraper {
 	return &StatSummaryScraper{
 		client:       client,
 		clusterCache: clusterCache,
+		nodeIndex:    newPersistedIndex[string]("node"),
+		pvcIndex:     newPersistedIndex[pvcKey]("pvc"),
 	}
 }
 
 func (s *StatSummaryScraper) Scrape() []metric.Update {
 
-	nodeNameToUID := buildNodeIndex(s.clusterCache.GetAllNodes())
-	pvcNameToUID := buildPVCIndex(s.clusterCache.GetAllPersistentVolumeClaims())
+	nodeNameToUID := s.nodeIndex.update(buildNodeIndex(s.clusterCache.GetAllNodes()))
+	pvcNameToUID := s.pvcIndex.update(buildPVCIndex(s.clusterCache.GetAllPersistentVolumeClaims()))
 
 	var scrapeResults []metric.Update
 	nodeStats, err := s.client.GetNodeData()