From 100f78d77166e5c377bd588445d2b072defb3a44 Mon Sep 17 00:00:00 2001
From: Alfonso Acosta <fons@weave.works>
Date: Tue, 10 May 2016 16:35:35 +0000
Subject: [PATCH] Provide truncation counts to the UI

---
 report/table.go      | 51 ++++++++++++++++++++++++--------------------
 report/table_test.go | 29 +++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 25 deletions(-)

diff --git a/report/table.go b/report/table.go
index 964da308dd..6b5346d908 100644
--- a/report/table.go
+++ b/report/table.go
@@ -1,57 +1,62 @@
 package report
 
 import (
+	"fmt"
 	"sort"
 	"strings"
 
+	log "github.com/Sirupsen/logrus"
 	"github.com/weaveworks/scope/common/mtime"
 )
 
 // MaxTableRows sets the limit on the table size to render
 // TODO: this won't be needed once we send reports incrementally
-const MaxTableRows = 20
+const (
+	MaxTableRows          = 20
+	TruncationCountPrefix = "table_truncation_count_"
+)
 
 // AddTable appends arbirary key-value pairs to the Node, returning a new node.
 func (node Node) AddTable(prefix string, labels map[string]string) Node {
 	count := 0
 	for key, value := range labels {
-		// It's enough to only include MaxTableRows+1
-		// since they won't be rendered anyhow
-		if count > MaxTableRows {
+		if count >= MaxTableRows {
 			break
 		}
 		node = node.WithLatest(prefix+key, mtime.Now(), value)
 		count++
-
+	}
+	if len(labels) > MaxTableRows {
+		truncationCount := fmt.Sprintf("%d", len(labels)-MaxTableRows)
+		node = node.WithLatest(TruncationCountPrefix+prefix, mtime.Now(), truncationCount)
 	}
 	return node
 }
 
 // ExtractTable returns the key-value pairs with the given prefix from this Node,
-func (node Node) ExtractTable(prefix string) (rows map[string]string, truncated bool) {
+func (node Node) ExtractTable(prefix string) (rows map[string]string, truncationCount int) {
 	rows = map[string]string{}
-	truncated = false
-	count := 0
+	truncationCount = 0
 	node.Latest.ForEach(func(key, value string) {
 		if strings.HasPrefix(key, prefix) {
-			if count >= MaxTableRows {
-				truncated = true
-				return
-			}
 			label := key[len(prefix):]
 			rows[label] = value
-			count++
 		}
 	})
-	return rows, truncated
+	if str, ok := node.Latest.Lookup(TruncationCountPrefix + prefix); ok {
+		if n, err := fmt.Sscanf(str, "%d", &truncationCount); n != 1 || err != nil {
+			log.Warn("Unexpected truncation count format %q", str)
+		}
+	}
+	return rows, truncationCount
 }
 
 // Table is the type for a table in the UI.
 type Table struct {
-	ID           string        `json:"id"`
-	Label        string        `json:"label"`
-	Rows         []MetadataRow `json:"rows"`
-	WasTruncated bool          `json:"was_truncated"`
+	ID              string        `json:"id"`
+	Label           string        `json:"label"`
+	Rows            []MetadataRow `json:"rows"`
+	TruncationCount int           `json:"truncation_count,omitempty"`
 }
 
 type tablesByID []Table
@@ -110,12 +115,12 @@ type TableTemplates map[string]TableTemplate
 func (t TableTemplates) Tables(node Node) []Table {
 	var result []Table
 	for _, template := range t {
-		rows, truncated := node.ExtractTable(template.Prefix)
+		rows, truncationCount := node.ExtractTable(template.Prefix)
 		table := Table{
-			ID:           template.ID,
-			Label:        template.Label,
-			Rows:         []MetadataRow{},
-			WasTruncated: truncated,
+			ID:              template.ID,
+			Label:           template.Label,
+			Rows:            []MetadataRow{},
+			TruncationCount: truncationCount,
 		}
 		keys := make([]string, 0, len(rows))
 		for k := range rows {
diff --git a/report/table_test.go b/report/table_test.go
index 7102995f88..cf276ba75c 100644
--- a/report/table_test.go
+++ b/report/table_test.go
@@ -1,6 +1,7 @@
 package report_test
 
 import (
+	"fmt"
 	"reflect"
 	"testing"
 
@@ -16,9 +17,9 @@ func TestTables(t *testing.T) {
 	nmd := report.MakeNode("foo1")
 
 	nmd = nmd.AddTable("foo_", want)
-	have, truncated := nmd.ExtractTable("foo_")
+	have, truncationCount := nmd.ExtractTable("foo_")
 
-	if truncated {
+	if truncationCount != 0 {
 		t.Error("Table shouldn't had been truncated")
 	}
 
@@ -26,3 +27,27 @@ func TestTables(t *testing.T) {
 		t.Error(test.Diff(want, have))
 	}
 }
+
+func TestTruncation(t *testing.T) {
+	wantTruncationCount := 1
+	want := map[string]string{}
+	for i := 0; i < report.MaxTableRows+wantTruncationCount; i++ {
+		key := fmt.Sprintf("key%d", i)
+		value := fmt.Sprintf("value%d", i)
+		want[key] = value
+	}
+
+	nmd := report.MakeNode("foo1")
+
+	nmd = nmd.AddTable("foo_", want)
+	_, truncationCount := nmd.ExtractTable("foo_")
+
+	if truncationCount != wantTruncationCount {
+		t.Error(
+			"Table should had been truncated by",
+			wantTruncationCount,
+			"and not",
+			truncationCount,
+		)
+	}
+}