Merge pull request #3269 from weaveworks/string-sets-add

Optimise WithParents() when there is only one parent

probe/awsecs/reporter.go

@@ -182,12 +182,12 @@ func (r Reporter) Tag(rpt report.Report) (report.Report, error) {
 
 			// parents sets to merge into all matching container nodes
 			parentsSets := report.MakeSets()
-			parentsSets = parentsSets.Add(report.ECSTask, report.MakeStringSet(taskID))
+			parentsSets = parentsSets.AddString(report.ECSTask, taskID)
 			if serviceName, ok := ecsInfo.TaskServiceMap[taskArn]; ok {
 				serviceID := report.MakeECSServiceNodeID(cluster, serviceName)
-				parentsSets = parentsSets.Add(report.ECSService, report.MakeStringSet(serviceID))
+				parentsSets = parentsSets.AddString(report.ECSService, serviceID)
 				// in addition, make service parent of task
-				rpt.ECSTask.Nodes[taskID] = rpt.ECSTask.Nodes[taskID].WithParents(report.MakeSets().Add(report.ECSService, report.MakeStringSet(serviceID)))
+				rpt.ECSTask.Nodes[taskID] = rpt.ECSTask.Nodes[taskID].WithParent(report.ECSService, serviceID)
 			}
 			for _, containerID := range info.ContainerIDs {
 				if containerNode, ok := rpt.Container.Nodes[containerID]; ok {

probe/docker/container.go

@@ -407,9 +407,7 @@ func (c *container) getBaseNode() report.Node {
 		ContainerCommand:  c.getSanitizedCommand(),
 		ImageID:           c.Image(),
 		ContainerHostname: c.Hostname(),
-	}).WithParents(report.MakeSets().
-		Add(report.ContainerImage, report.MakeStringSet(report.MakeContainerImageNodeID(c.Image()))),
-	)
+	}).WithParent(report.ContainerImage, report.MakeContainerImageNodeID(c.Image()))
 	result = result.AddPrefixPropertyList(LabelPrefix, c.container.Config.Labels)
 	if !c.noEnvironmentVariables {
 		result = result.AddPrefixPropertyList(EnvPrefix, c.env())

probe/docker/tagger.go

@@ -74,7 +74,7 @@ func (t *Tagger) Tag(r report.Report) (report.Report, error) {
 		})
 		r.SwarmService.AddNode(node)
 
-		r.Container.Nodes[containerID] = container.WithParents(container.Parents.Add(report.SwarmService, report.MakeStringSet(nodeID)))
+		r.Container.Nodes[containerID] = container.WithParent(report.SwarmService, nodeID)
 	}
 
 	return r, nil
@@ -116,17 +116,13 @@ func (t *Tagger) tag(tree process.Tree, topology *report.Topology) {
 		}
 
 		node = node.WithLatest(ContainerID, mtime.Now(), c.ID())
-		node = node.WithParents(report.MakeSets().
-			Add(report.Container, report.MakeStringSet(report.MakeContainerNodeID(c.ID()))),
-		)
+		node = node.WithParent(report.Container, report.MakeContainerNodeID(c.ID()))
 
 		// If we can work out the image name, add a parent tag for it
 		image, ok := t.registry.GetContainerImage(c.Image())
 		if ok && len(image.RepoTags) > 0 {
 			imageName := ImageNameWithoutTag(image.RepoTags[0])
-			node = node.WithParents(report.MakeSets().
-				Add(report.ContainerImage, report.MakeStringSet(report.MakeContainerImageNodeID(imageName))),
-			)
+			node = node.WithParent(report.ContainerImage, report.MakeContainerImageNodeID(imageName))
 		}
 
 		topology.ReplaceNode(node)

probe/host/tagger.go

@@ -26,14 +26,13 @@ func (Tagger) Name() string { return "Host" }
 func (t Tagger) Tag(r report.Report) (report.Report, error) {
 	var (
 		metadata = map[string]string{report.HostNodeID: t.hostNodeID}
-		parents  = report.MakeSets().Add(report.Host, report.MakeStringSet(t.hostNodeID))
 	)
 
 	// Explicitly don't tag Endpoints, Addresses and Overlay nodes - These topologies include pseudo nodes,
 	// and as such do their own host tagging.
 	for _, topology := range []report.Topology{r.Process, r.Container, r.ContainerImage, r.Host, r.Pod} {
 		for _, node := range topology.Nodes {
-			topology.ReplaceNode(node.WithLatests(metadata).WithParents(parents))
+			topology.ReplaceNode(node.WithLatests(metadata).WithParent(report.Host, t.hostNodeID))
 		}
 	}
 	return r, nil

probe/kubernetes/pod.go

@@ -56,7 +56,7 @@ func (p *pod) UID() string {
 }
 
 func (p *pod) AddParent(topology, id string) {
-	p.parents = p.parents.Add(topology, report.MakeStringSet(id))
+	p.parents = p.parents.AddString(topology, id)
 }
 
 func (p *pod) State() string {

probe/kubernetes/reporter.go

@@ -249,10 +249,7 @@ func (r *Reporter) Tag(rpt report.Report) (report.Report, error) {
 			n = n.WithLatest(report.DoesNotMakeConnections, mtime.Now(), "")
 		}
 
-		rpt.Container.Nodes[id] = n.WithParents(report.MakeSets().Add(
-			report.Pod,
-			report.MakeStringSet(report.MakePodNodeID(uid)),
-		))
+		rpt.Container.Nodes[id] = n.WithParent(report.Pod, report.MakePodNodeID(uid))
 	}
 	return rpt, nil
 }
@@ -357,7 +354,7 @@ func (r *Reporter) hostTopology(services []Service) report.Topology {
 	t := report.MakeTopology()
 	t.AddNode(
 		report.MakeNode(report.MakeHostNodeID(r.hostID)).
-			WithSets(report.MakeSets().Add(host.LocalNetworks, report.MakeStringSet(serviceNetwork.String()))))
+			WithSets(report.MakeSets().AddString(host.LocalNetworks, serviceNetwork.String())))
 	return t
 }
 

probe/kubernetes/reporter_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	apiv1 "k8s.io/api/core/v1"
+	apiv1beta1 "k8s.io/api/extensions/v1beta1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 
@@ -111,9 +112,10 @@ func newMockClient() *mockClient {
 }
 
 type mockClient struct {
-	pods     []kubernetes.Pod
-	services []kubernetes.Service
-	logs     map[string]io.ReadCloser
+	pods        []kubernetes.Pod
+	services    []kubernetes.Service
+	deployments []kubernetes.Deployment
+	logs        map[string]io.ReadCloser
 }
 
 func (c *mockClient) Stop() {}
@@ -143,6 +145,11 @@ func (c *mockClient) WalkCronJobs(f func(kubernetes.CronJob) error) error {
 	return nil
 }
 func (c *mockClient) WalkDeployments(f func(kubernetes.Deployment) error) error {
+	for _, deployment := range c.deployments {
+		if err := f(deployment); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 func (c *mockClient) WalkNamespaces(f func(kubernetes.NamespaceResource) error) error {
@@ -282,6 +289,39 @@ func TestReporter(t *testing.T) {
 
 }
 
+func BenchmarkReporter(b *testing.B) {
+	hr := controls.NewDefaultHandlerRegistry()
+	mockK8s := newMockClient()
+	// Add more dummy data
+	for i := 0; i < 50; i++ {
+		service := apiService1
+		service.ObjectMeta.UID = types.UID(fmt.Sprintf("service%d", i))
+		mockK8s.services = append(mockK8s.services, kubernetes.NewService(&service))
+		pod := apiPod1
+		pod.ObjectMeta.UID = types.UID(fmt.Sprintf("pod%d", i))
+		mockK8s.pods = append(mockK8s.pods, kubernetes.NewPod(&pod))
+		deployment := apiv1beta1.Deployment{
+			TypeMeta: metav1.TypeMeta{
+				Kind:       "Deployment",
+				APIVersion: "v1beta1",
+			},
+			ObjectMeta: metav1.ObjectMeta{
+				Name:              fmt.Sprintf("deployment%d", i),
+				UID:               types.UID(fmt.Sprintf("deployment%d", i)),
+				Namespace:         "ping",
+				CreationTimestamp: metav1.Now(),
+			},
+		}
+		mockK8s.deployments = append(mockK8s.deployments, kubernetes.NewDeployment(&deployment))
+	}
+	reporter := kubernetes.NewReporter(mockK8s, nil, "probe-id", "foo", nil, hr, nodeName, 0)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		reporter.Report()
+	}
+}
+
 func TestTagger(t *testing.T) {
 	rpt := report.MakeReport()
 	rpt.Container.AddNode(report.MakeNodeWith("container1", map[string]string{

probe/overlay/weave.go

@@ -330,7 +330,7 @@ func (w *Weave) addCurrentPeerInfo(latests map[string]string, node report.Node)
 		latests[WeavePluginDriver] = w.statusCache.Plugin.DriverName
 	}
 	node = node.AddPrefixMulticolumnTable(WeaveConnectionsMulticolumnTablePrefix, getConnectionsTable(w.statusCache.Router))
-	node = node.WithParents(report.MakeSets().Add(report.Host, report.MakeStringSet(w.hostID)))
+	node = node.WithParent(report.Host, w.hostID)
 
 	return latests, node
 }

render/container.go

@@ -150,7 +150,7 @@ func (r containerWithImageNameRenderer) Render(rpt report.Report) Nodes {
 		c = propagateLatest(docker.ImageLabelPrefix+"works.weave.role", image, c)
 		c.Parents = c.Parents.
 			Delete(report.ContainerImage).
-			Add(report.ContainerImage, report.MakeStringSet(imageNodeID))
+			AddString(report.ContainerImage, imageNodeID)
 		outputs[id] = c
 	}
 	return Nodes{Nodes: outputs, Filtered: containers.Filtered}

report/node.go

@@ -152,6 +152,12 @@ func (n Node) WithLatestControl(control string, ts time.Time, data NodeControlDa
 	return n
 }
 
+// WithParent returns a fresh copy of n, with one parent added
+func (n Node) WithParent(key, parent string) Node {
+	n.Parents = n.Parents.AddString(key, parent)
+	return n
+}
+
 // WithParents returns a fresh copy of n, with sets merged in.
 func (n Node) WithParents(parents Sets) Node {
 	n.Parents = n.Parents.Merge(parents)

report/sets.go

@@ -46,6 +46,21 @@ func (s Sets) Add(key string, value StringSet) Sets {
 	}
 }
 
+// AddString adds a single string under a key, creating a new StringSet if necessary.
+func (s Sets) AddString(key string, str string) Sets {
+	if s.psMap == nil {
+		s = emptySets
+	}
+	value, found := s.Lookup(key)
+	if found && value.Contains(str) {
+		return s
+	}
+	value = value.Add(str)
+	return Sets{
+		psMap: s.psMap.Set(key, value),
+	}
+}
+
 // Delete the given set from the Sets.
 func (s Sets) Delete(key string) Sets {
 	if s.psMap == nil {