From 05b350850f72fce4d82fea70ab74d74c63738185 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Thu, 11 Oct 2018 13:53:56 +0000
Subject: [PATCH] Drop topologies which are way too big

---
 app/multitenant/aws_collector.go |  4 ++++
 prog/app.go                      |  5 +++--
 prog/main.go                     |  2 ++
 report/report.go                 | 12 ++++++++++++
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/app/multitenant/aws_collector.go b/app/multitenant/aws_collector.go
index 20415d123..e1d78b1d9 100644
--- a/app/multitenant/aws_collector.go
+++ b/app/multitenant/aws_collector.go
@@ -104,6 +104,7 @@ type AWSCollectorConfig struct {
 	NatsHost       string
 	MemcacheClient *MemcacheClient
 	Window         time.Duration
+	MaxTopNodes    int
 }
 
 type awsCollector struct {
@@ -312,6 +313,9 @@ func (c *awsCollector) getReports(ctx context.Context, reportKeys []string) ([]r
 			log.Warningf("Error fetching from cache: %v", err)
 		}
 		for key, report := range found {
+			if c.cfg.MaxTopNodes > 0 {
+				report = report.DropTopologiesOver(c.cfg.MaxTopNodes)
+			}
 			report = report.Upgrade()
 			c.inProcess.StoreReport(key, report)
 			reports = append(reports, report)
diff --git a/prog/app.go b/prog/app.go
index 234fe51e1..00d9568e4 100644
--- a/prog/app.go
+++ b/prog/app.go
@@ -82,7 +82,7 @@ func router(collector app.Collector, controlRouter app.ControlRouter, pipeRouter
 }
 
 func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL, natsHostname string,
-	memcacheConfig multitenant.MemcacheConfig, window time.Duration, createTables bool) (app.Collector, error) {
+	memcacheConfig multitenant.MemcacheConfig, window time.Duration, maxTopNodes int, createTables bool) (app.Collector, error) {
 	if collectorURL == "local" {
 		return app.NewCollector(window), nil
 	}
@@ -124,6 +124,7 @@ func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL, natsHo
 				NatsHost:       natsHostname,
 				MemcacheClient: memcacheClient,
 				Window:         window,
+				MaxTopNodes:    maxTopNodes,
 			},
 		)
 		if err != nil {
@@ -232,7 +233,7 @@ func appMain(flags appFlags) {
 			Service:          flags.memcachedService,
 			CompressionLevel: flags.memcachedCompressionLevel,
 		},
-		flags.window, flags.awsCreateTables)
+		flags.window, flags.maxTopNodes, flags.awsCreateTables)
 	if err != nil {
 		log.Fatalf("Error creating collector: %v", err)
 		return
diff --git a/prog/main.go b/prog/main.go
index 9630b816b..b20f02cd1 100644
--- a/prog/main.go
+++ b/prog/main.go
@@ -140,6 +140,7 @@ type probeFlags struct {
 
 type appFlags struct {
 	window         time.Duration
+	maxTopNodes    int
 	listen         string
 	stopTimeout    time.Duration
 	logLevel       string
@@ -342,6 +343,7 @@ func setupFlags(flags *flags) {
 
 	// App flags
 	flag.DurationVar(&flags.app.window, "app.window", 15*time.Second, "window")
+	flag.IntVar(&flags.app.maxTopNodes, "app.max-topology-nodes", 10000, "drop topologies with more than this many nodes (0 to disable)")
 	flag.StringVar(&flags.app.listen, "app.http.address", ":"+strconv.Itoa(xfer.AppPort), "webserver listen address")
 	flag.DurationVar(&flags.app.stopTimeout, "app.stopTimeout", 5*time.Second, "How long to wait for http requests to finish when shutting down")
 	flag.StringVar(&flags.app.logLevel, "app.log.level", "info", "logging threshold level: debug|info|warn|error|fatal|panic")
diff --git a/report/report.go b/report/report.go
index f40643d9f..1ee60da91 100644
--- a/report/report.go
+++ b/report/report.go
@@ -417,6 +417,18 @@ func (r Report) Validate() error {
 	return nil
 }
 
+// DropTopologiesOver - as a protection against overloading the app
+// server, drop topologies that have really large node counts. In
+// practice we only see this with runaway numbers of zombie processes.
+func (r Report) DropTopologiesOver(limit int) Report {
+	r.WalkNamedTopologies(func(name string, topology *Topology) {
+		if topology != nil && len(topology.Nodes) > limit {
+			topology.Nodes = Nodes{}
+		}
+	})
+	return r
+}
+
 // Upgrade returns a new report based on a report received from the old probe.
 //
 func (r Report) Upgrade() Report {