Merge pull request #1257 from kubecost/sean/recon-dedup

get reconciliation map

pkg/kubecost/asset.go

@@ -2652,6 +2652,49 @@ func (as *AssetSet) ReconciliationMatch(query Asset) (Asset, bool, error) {
 	return nil, false, fmt.Errorf("Asset not found to match %s", query)
 }
 
+// ReconciliationMatchMap returns a map of the calling AssetSet's Assets, by provider id and category. This data structure
+// allows for reconciliation matching to be done in constant time and prevents duplicate reconciliation.
+func (as *AssetSet) ReconciliationMatchMap() map[string]map[string]Asset {
+	matchMap := make(map[string]map[string]Asset)
+
+	if as == nil {
+		return matchMap
+	}
+
+	for _, asset := range as.assets {
+		if asset == nil {
+			continue
+		}
+		props := asset.Properties()
+		// Ignore cloud assets and assets that cannot be matched when looking for reconciliation matches
+		if props == nil || props.ProviderID == "" || asset.Type() == CloudAssetType {
+			continue
+		}
+
+		if _, ok := matchMap[props.ProviderID]; !ok {
+			matchMap[props.ProviderID] = make(map[string]Asset)
+		}
+
+		// Check if a match is already in the map
+		if duplicateAsset, ok := matchMap[props.ProviderID][props.Category]; ok {
+			log.DedupedWarningf(5, "duplicate asset found when reconciling for %s", props.ProviderID)
+			// if one asset already has adjustment use that one
+			if duplicateAsset.Adjustment() == 0 && asset.Adjustment() != 0 {
+				matchMap[props.ProviderID][props.Category] = asset
+			} else if duplicateAsset.Adjustment() != 0 && asset.Adjustment() == 0 {
+				matchMap[props.ProviderID][props.Category] = duplicateAsset
+				// otherwise use the one with the higher cost
+			} else if duplicateAsset.TotalCost() < asset.TotalCost() {
+				matchMap[props.ProviderID][props.Category] = asset
+			}
+		} else {
+			matchMap[props.ProviderID][props.Category] = asset
+		}
+
+	}
+	return matchMap
+}
+
 // Get returns the Asset in the AssetSet at the given key, or nil and false
 // if no Asset exists for the given key
 func (as *AssetSet) Get(key string) (Asset, bool) {
@@ -2838,25 +2881,24 @@ func (as *AssetSet) accumulate(that *AssetSet) (*AssetSet, error) {
 	return acc, nil
 }
 
-
 type DiffKind string
 
 const (
-    DiffAdded DiffKind = "added"
-    DiffRemoved = "removed"
+	DiffAdded   DiffKind = "added"
+	DiffRemoved          = "removed"
 )
 
 // Diff stores an object and a string that denotes whether that object was
 // added or removed from a set of those objects
 type Diff[T any] struct {
 	Entity T
-	Kind DiffKind
+	Kind   DiffKind
 }
 
 // DiffAsset takes two AssetSets and returns a slice of Diffs by checking
 // the keys of each AssetSet. If a key is not found, a Diff is generated
 // and added to the slice.
-func DiffAsset(before, after *AssetSet) []Diff[Asset]{
+func DiffAsset(before, after *AssetSet) []Diff[Asset] {
 	changedItems := []Diff[Asset]{}
 
 	for assetKey1, asset1 := range before.assets {
@@ -2865,7 +2907,7 @@ func DiffAsset(before, after *AssetSet) []Diff[Asset]{
 			changedItems = append(changedItems, d)
 		}
 	}
-	
+
 	for assetKey2, asset2 := range after.assets {
 		if _, ok := before.assets[assetKey2]; !ok {
 			d := Diff[Asset]{asset2, DiffAdded}

pkg/kubecost/asset_test.go

@@ -850,6 +850,32 @@ func TestAssetSet_InsertMatchingWindow(t *testing.T) {
 	})
 }
 
+func TestAssetSet_ReconciliationMatchMap(t *testing.T) {
+	endYesterday := time.Now().UTC().Truncate(day)
+	startYesterday := endYesterday.Add(-day)
+
+	as := GenerateMockAssetSet(startYesterday)
+	matchMap := as.ReconciliationMatchMap()
+
+	// Determine the number of assets by provider ID
+	assetCountByProviderId := make(map[string]int, len(matchMap))
+	as.Each(func(key string, a Asset) {
+		if a == nil || a.Properties() == nil || a.Properties().ProviderID == "" {
+			return
+		}
+		if _, ok := assetCountByProviderId[a.Properties().ProviderID]; !ok {
+			assetCountByProviderId[a.Properties().ProviderID] = 0
+		}
+		assetCountByProviderId[a.Properties().ProviderID] += 1
+	})
+
+	for k, count := range assetCountByProviderId {
+		if len(matchMap[k]) != count {
+			t.Errorf("AssetSet.ReconciliationMatchMap: incorrect asset count for provider id: %s", k)
+		}
+	}
+}
+
 func TestAssetSetRange_Accumulate(t *testing.T) {
 	endYesterday := time.Now().UTC().Truncate(day)
 	startYesterday := endYesterday.Add(-day)