Fix index creation. Batch enrich

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

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

@@ -36,20 +36,21 @@ func newDCGMTargetScraper(provider target.TargetProvider, enrich UpdateEnricher)
 // resolved against a freshly built index of the cluster's current pods. Left unset if pod_uid
 // is already present, or if namespace/pod can't be resolved to a known pod.
 func podUIDEnricher(clusterCache clustercache.ClusterCache) UpdateEnricher {
-	index := buildPodIndex(clusterCache.GetAllPods())
-	return func(update metric.Update) metric.Update {
-		if update.Labels[source.PodUIDLabel] != "" {
-			return update
-		}
-		namespace, pod := update.Labels[source.NamespaceLabel], update.Labels[source.PodLabel]
-		if namespace == "" || pod == "" {
-			return update
-		}
+	return func(updates []metric.Update) {
+		index := buildPodIndex(clusterCache.GetAllPods())
+		for _, update := range updates {
+			if update.Labels[source.PodUIDLabel] != "" {
+				continue
+			}
+			namespace, pod := update.Labels[source.NamespaceLabel], update.Labels[source.PodLabel]
+			if namespace == "" || pod == "" {
+				continue
+			}
 
-		if uid, ok := index[podKey{namespace: namespace, name: pod}]; ok {
-			update.Labels[source.PodUIDLabel] = string(uid)
+			if uid, ok := index[podKey{namespace: namespace, name: pod}]; ok {
+				update.Labels[source.PodUIDLabel] = string(uid)
+			}
 		}
-		return update
 	}
 }
 

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

@@ -71,84 +71,144 @@ func Test_podUIDEnricher(t *testing.T) {
 	cache := &clustercache.MockClusterCache{
 		Pods: []*clustercache.Pod{
 			{UID: "pod-uid-1", Name: "pod1", Namespace: "namespace1"},
+			{UID: "pod-uid-2", Name: "pod2", Namespace: "namespace2"},
 		},
 	}
-	enrich := podUIDEnricher(cache)
 
 	tests := map[string]struct {
-		update metric.Update
-		want   metric.Update
+		labels map[string]string
+		want   map[string]string
 	}{
 		"resolves pod_uid from namespace and pod": {
-			update: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "pod1",
-				},
+			labels: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "pod1",
 			},
-			want: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "pod1",
-					source.PodUIDLabel:    "pod-uid-1",
-				},
+			want: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "pod1",
+				source.PodUIDLabel:    "pod-uid-1",
 			},
 		},
 		"unknown pod is left unset": {
-			update: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "unknown-pod",
-				},
+			labels: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "unknown-pod",
 			},
-			want: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "unknown-pod",
-				},
+			want: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "unknown-pod",
 			},
 		},
-		"missing namespace or pod label is left unset": {
-			update: metric.Update{
-				Labels: map[string]string{
-					source.PodLabel: "pod1",
-				},
+		"pod name from wrong namespace is left unset": {
+			labels: map[string]string{
+				source.NamespaceLabel: "namespace2",
+				source.PodLabel:       "pod1",
 			},
-			want: metric.Update{
-				Labels: map[string]string{
-					source.PodLabel: "pod1",
-				},
+			want: map[string]string{
+				source.NamespaceLabel: "namespace2",
+				source.PodLabel:       "pod1",
 			},
 		},
+		"missing namespace label is left unset": {
+			labels: map[string]string{
+				source.PodLabel: "pod1",
+			},
+			want: map[string]string{
+				source.PodLabel: "pod1",
+			},
+		},
+		"missing pod label is left unset": {
+			labels: map[string]string{
+				source.NamespaceLabel: "namespace1",
+			},
+			want: map[string]string{
+				source.NamespaceLabel: "namespace1",
+			},
+		},
+		"nil labels is left untouched": {
+			labels: nil,
+			want:   nil,
+		},
 		"existing non-empty pod_uid is left untouched": {
-			update: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "pod1",
-					source.PodUIDLabel:    "already-set",
-				},
+			labels: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "pod1",
+				source.PodUIDLabel:    "already-set",
 			},
-			want: metric.Update{
-				Labels: map[string]string{
-					source.NamespaceLabel: "namespace1",
-					source.PodLabel:       "pod1",
-					source.PodUIDLabel:    "already-set",
-				},
+			want: map[string]string{
+				source.NamespaceLabel: "namespace1",
+				source.PodLabel:       "pod1",
+				source.PodUIDLabel:    "already-set",
 			},
 		},
 	}
 
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
-			got := enrich(tt.update)
-			if len(got.Labels) != len(tt.want.Labels) {
-				t.Fatalf("got labels %v, want %v", got.Labels, tt.want.Labels)
-			}
-			for k, v := range tt.want.Labels {
-				if got.Labels[k] != v {
-					t.Errorf("label %q = %q, want %q", k, got.Labels[k], v)
-				}
-			}
+			enrich := podUIDEnricher(cache)
+			updates := []metric.Update{{Labels: cloneLabels(tt.labels)}}
+
+			enrich(updates)
+
+			assertLabels(t, updates[0].Labels, tt.want)
+		})
+	}
+
+	// A single call must correctly enrich every update in the batch, not just
+	// the first, since updates are processed together rather than one at a time.
+	t.Run("enriches every update in a multi-update batch", func(t *testing.T) {
+		enrich := podUIDEnricher(cache)
+		updates := []metric.Update{
+			{Labels: map[string]string{source.NamespaceLabel: "namespace1", source.PodLabel: "pod1"}},
+			{Labels: map[string]string{source.NamespaceLabel: "namespace2", source.PodLabel: "pod2"}},
+			{Labels: map[string]string{source.NamespaceLabel: "namespace1", source.PodLabel: "unknown-pod"}},
+		}
+
+		enrich(updates)
+
+		assertLabels(t, updates[0].Labels, map[string]string{
+			source.NamespaceLabel: "namespace1",
+			source.PodLabel:       "pod1",
+			source.PodUIDLabel:    "pod-uid-1",
+		})
+		assertLabels(t, updates[1].Labels, map[string]string{
+			source.NamespaceLabel: "namespace2",
+			source.PodLabel:       "pod2",
+			source.PodUIDLabel:    "pod-uid-2",
+		})
+		assertLabels(t, updates[2].Labels, map[string]string{
+			source.NamespaceLabel: "namespace1",
+			source.PodLabel:       "unknown-pod",
 		})
+	})
+
+	t.Run("empty batch does not panic", func(t *testing.T) {
+		enrich := podUIDEnricher(cache)
+		enrich(nil)
+		enrich([]metric.Update{})
+	})
+}
+
+func cloneLabels(labels map[string]string) map[string]string {
+	if labels == nil {
+		return nil
+	}
+	clone := make(map[string]string, len(labels))
+	for k, v := range labels {
+		clone[k] = v
+	}
+	return clone
+}
+
+func assertLabels(t *testing.T, got, want map[string]string) {
+	t.Helper()
+	if len(got) != len(want) {
+		t.Fatalf("got labels %v, want %v", got, want)
+	}
+	for k, v := range want {
+		if got[k] != v {
+			t.Errorf("label %q = %q, want %q", k, got[k], v)
+		}
 	}
 }

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

@@ -12,9 +12,9 @@ import (
 	"github.com/opencost/opencost/modules/collector-source/pkg/scrape/target"
 )
 
-// UpdateEnricher optionally transforms a single scraped update before it's returned from
+// UpdateEnricher optionally batch transforms entire set of updates before they returned from
 // Scrape().
-type UpdateEnricher func(update metric.Update) metric.Update
+type UpdateEnricher func(update []metric.Update)
 
 type TargetScraper struct {
 	name           string // identifier for the scraper
@@ -81,9 +81,7 @@ func (s *TargetScraper) Scrape() []metric.Update {
 					Labels: result.Labels,
 					Value:  result.Value,
 				}
-				if s.enrich != nil {
-					update = s.enrich(update)
-				}
+
 				scrapeResults = append(scrapeResults, update)
 			}
 			return scrapeResults
@@ -93,6 +91,10 @@ func (s *TargetScraper) Scrape() []metric.Update {
 
 	updates := concurrentScrape(scrapeFuncs...)
 
+	if s.enrich != nil {
+		s.enrich(updates)
+	}
+
 	// dispatch a scrape event for this specific scrape
 	events.Dispatch(event.ScrapeEvent{
 		ScraperName: s.name,

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

@@ -454,12 +454,13 @@ func TestTargetScraper_Scrape(t *testing.T) {
 			name:       "Enrichment",
 			scrapeText: enrichScrape,
 			targetScraperFactory: func(provider target.TargetProvider) *TargetScraper {
-				enrich := func(update metric.Update) metric.Update {
-					if update.Labels["extra"] != "" {
-						return update
+				enrich := func(updates []metric.Update) {
+					for i := range updates {
+						if updates[i].Labels["extra"] != "" {
+							continue
+						}
+						updates[i].Labels["extra"] = "enriched"
 					}
-					update.Labels["extra"] = "enriched"
-					return update
 				}
 				return newTargetScrapper("test-enrich", provider, nil, false, enrich)
 			},