From 7618b368ac3ad75daefd86cff5f6b183d2e5bafb Mon Sep 17 00:00:00 2001
From: Roberto Bruggemann <roberto@weave.works>
Date: Fri, 8 Dec 2017 17:44:18 +0000
Subject: [PATCH] report.Upgrade() add deployments to pods as parent

This was previously done selectPodsWithDeployments.Render().
---
 render/pod.go    | 29 +----------------------------
 report/report.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 28 deletions(-)

diff --git a/render/pod.go b/render/pod.go
index a9ee1e73d1..9180db95f1 100644
--- a/render/pod.go
+++ b/render/pod.go
@@ -62,9 +62,7 @@ var PodRenderer = Memoise(ConditionalRenderer(renderKubernetesTopologies,
 					),
 				),
 			),
-			// ConnectionJoin invokes the renderer twice, hence it
-			// helps to memoise it.
-			ConnectionJoin(MapPod2IP, Memoise(selectPodsWithDeployments{})),
+			ConnectionJoin(MapPod2IP, SelectPod),
 		),
 	),
 ))
@@ -113,31 +111,6 @@ func renderParents(childTopology string, parentTopologies []string, noParentsPse
 	)...)
 }
 
-// Renderer to return modified Pod nodes to elide replica sets and point directly
-// to deployments where applicable.
-type selectPodsWithDeployments struct{}
-
-func (s selectPodsWithDeployments) Render(rpt report.Report) Nodes {
-	result := report.Nodes{}
-	// For each pod, we check for any replica sets, and merge any deployments they point to
-	// into a replacement Parents value.
-	for podID, pod := range rpt.Pod.Nodes {
-		if replicaSetIDs, ok := pod.Parents.Lookup(report.ReplicaSet); ok {
-			newParents := pod.Parents.Delete(report.ReplicaSet)
-			for _, replicaSetID := range replicaSetIDs {
-				if replicaSet, ok := rpt.ReplicaSet.Nodes[replicaSetID]; ok {
-					if deploymentIDs, ok := replicaSet.Parents.Lookup(report.Deployment); ok {
-						newParents = newParents.Add(report.Deployment, deploymentIDs)
-					}
-				}
-			}
-			pod = pod.WithParents(newParents)
-		}
-		result[podID] = pod
-	}
-	return Nodes{Nodes: result}
-}
-
 // MapPod2IP maps pod nodes to their IP address.  This allows pods to
 // be joined directly with the endpoint topology.
 func MapPod2IP(m report.Node) []string {
diff --git a/report/report.go b/report/report.go
index edb52e4239..fe5e205b38 100644
--- a/report/report.go
+++ b/report/report.go
@@ -341,6 +341,10 @@ func (r Report) Validate() error {
 //
 // This for now creates node's LatestControls from Controls.
 func (r Report) Upgrade() Report {
+	return r.upgradeLatestControls().upgradePodNodes()
+}
+
+func (r Report) upgradeLatestControls() Report {
 	needUpgrade := false
 	r.WalkTopologies(func(topology *Topology) {
 		for _, node := range topology.Nodes {
@@ -349,6 +353,7 @@ func (r Report) Upgrade() Report {
 			}
 		}
 	})
+
 	if !needUpgrade {
 		return r
 	}
@@ -369,9 +374,47 @@ func (r Report) Upgrade() Report {
 		}
 		topology.Nodes = n
 	})
+
 	return cp
 }
 
+func (r Report) upgradePodNodes() Report {
+	needUpgrade := false
+	for _, pod := range r.Pod.Nodes {
+		_, hasReplicaSets := pod.Parents.Lookup(ReplicaSet)
+		// ReplicaSets are not reported since adding Deployments as pods' parent
+		if hasReplicaSets {
+			needUpgrade = true
+			break
+		}
+	}
+
+	if !needUpgrade {
+		return r
+	}
+
+	// For each pod, we check for any replica sets, and merge any deployments they point to
+	// into a replacement Parents value.
+	nodes := Nodes{}
+	for podID, pod := range r.Pod.Nodes {
+		if replicaSetIDs, ok := pod.Parents.Lookup(ReplicaSet); ok {
+			newParents := pod.Parents.Delete(ReplicaSet)
+			for _, replicaSetID := range replicaSetIDs {
+				if replicaSet, ok := r.ReplicaSet.Nodes[replicaSetID]; ok {
+					if deploymentIDs, ok := replicaSet.Parents.Lookup(Deployment); ok {
+						newParents = newParents.Add(Deployment, deploymentIDs)
+					}
+				}
+			}
+			pod = pod.WithParents(newParents)
+		}
+		nodes[podID] = pod
+	}
+	r.Pod.Nodes = nodes
+
+	return r
+}
+
 // BackwardCompatible returns a new backward-compatible report.
 //
 // This for now creates node's Controls from LatestControls.