logdog: Add bundler library.

client/internal/logdog/butler/bundler/bundler_test.go

Index: client/internal/logdog/butler/bundler/bundler_test.go
diff --git a/client/internal/logdog/butler/bundler/bundler_test.go b/client/internal/logdog/butler/bundler/bundler_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..1b66729f5a63ab9ce7f2967c05ee11e68b4ba477
--- /dev/null
+++ b/client/internal/logdog/butler/bundler/bundler_test.go
@@ -0,0 +1,394 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package bundler
+
+import (
+	"crypto/md5"
+	"encoding/hex"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/luci/luci-go/common/logdog/protocol"
+	"github.com/luci/luci-go/common/logdog/protocol/protoutil"
+	"github.com/luci/luci-go/common/logdog/types"
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+var (
+	testNow = time.Date(2015, 1, 1, 0, 0, 0, 0, time.UTC)
+)
+
+// fakeSizer is a Sizer implementation that counts (obviously incorrect) fixed
+// sizes for each entry type.
+type fakeSizer struct {
+	Bundle      int64
+	BundleEntry int64
+	LogEntry    int64
+
+	size int64
+	seen map[types.StreamPath]bool
+
+	lastSize int64
+	lastSeen types.StreamPath
+}
+
+func (s *fakeSizer) Size() int64 {
+	return s.Bundle + s.size
+}
+
+func (s *fakeSizer) Append(be *protocol.ButlerLogBundle_Entry, e *protocol.LogEntry) {
+	size := int64(0)
+
+	// Add the ButlerLogBundle_Entry code if we haven't seen it before.
+	path := protoutil.DescriptorPath(be.GetDesc())
+	if seen := s.seen[path]; !seen {
+		if s.seen == nil {
+			s.seen = map[types.StreamPath]bool{
+				path: true,
+			}
+		} else {
+			s.seen[path] = true
+		}
+		s.lastSeen = path
+		size += s.BundleEntry
+	} else {
+		s.lastSeen = ""
+	}
+
+	if e != nil {
+		// Each character in the line gets LogEntry space.
+		if len(e.GetLines()) == 1 {
+			size += s.LogEntry * int64(len(e.GetLines()[0]))
+		} else {
+			size += s.LogEntry
+		}
+	}
+
+	s.lastSize = size
+	s.size += size
+}
+
+func (s *fakeSizer) Undo() {
+	s.size -= s.lastSize
+	if s.lastSeen != "" {
+		s.seen[s.lastSeen] = false
+	}
+}
+
+func hash(s, t string) []byte {
+	sum := md5.Sum([]byte(fmt.Sprintf("%s::%s", s, t)))
+	return sum[:]
+}
+
+func key(s, t string) string {
+	return hex.EncodeToString(hash(s, t))
+}
+
+// addEntry generates a ButlerLogBundle_Entry and appends it to our Bundler via
+// one or more calls to Append.
+//
+// If "le" strings are supplied, those will create generated LogEntry for that
+// ButlerLogBundle_Entry.
+func gen(e string, t bool, le ...string) *protocol.ButlerLogBundle_Entry {
+	name := key(e, "name")
+	contentType := "test/data"
+
+	be := &protocol.ButlerLogBundle_Entry{
+		Desc: &protocol.LogStreamDescriptor{
+			Prefix:      &e,
+			Name:        &name,
+			ContentType: &contentType,
+			Timestamp:   protoutil.NewTimestamp(testNow),
+		},
+		Terminal: &t,
+	}
+
+	if len(le) > 0 {
+		be.Logs = make([]*protocol.LogEntry, len(le))
+		for i, l := range le {
+			be.Logs[i] = &protocol.LogEntry{
+				Lines: []string{
+					l,
+				},
+				Data: []*protocol.LogEntry_Data{
+					{
+						Value: hash(l, "data0"),
+					},
+					{
+						Value: hash(l, "data1"),
+					},
+					{
+						Value: hash(l, "data2"),
+					},
+				},
+			}
+		}
+	}
+	return be
+}
+
+func logEntryName(le *protocol.LogEntry) string {
+	if len(le.GetLines()) != 1 {
+		return ""
+	}
+	return le.GetLines()[0]
+}
+
+// "expected" is a notation to express a bundle entry and its keys:
+//   "a": a bundle entry keyed on "a".
+//   "+a": a terminal bundle entry keyed on "a".
+//   "a:1:2:3": a bundle entry keyed on "a" with three log entries, each keyed on
+//            "1", "2", and "3" respectively.
+func shouldHaveBundleEntries(actual interface{}, expected ...interface{}) string {
+	bundle := actual.(*protocol.ButlerLogBundle)
+
+	errors := []string{}
+	fail := func(f string, args ...interface{}) {
+		errors = append(errors, fmt.Sprintf(f, args...))
+	}
+
+	term := make(map[string]bool)
+	exp := make(map[string][]string)
+
+	// Parse expectation strings.
+	for _, e := range expected {
+		s := e.(string)
+		if len(s) == 0 {
+			continue
+		}
+
+		t := false
+		if s[0] == '+' {
+			t = true
+			s = s[1:]
+		}
+
+		parts := strings.Split(s, ":")
+		name := parts[0]
+		term[name] = t
+
+		if len(parts) > 1 {
+			exp[name] = append(exp[name], parts[1:]...)
+		}
+	}
+
+	entries := make(map[string]*protocol.ButlerLogBundle_Entry)
+	for _, be := range bundle.GetEntries() {
+		entries[be.GetDesc().GetPrefix()] = be
+	}
+	for name, t := range term {
+		be := entries[name]
+		if be == nil {
+			fail("No bundle entry for [%s]", name)
+			continue
+		}
+		delete(entries, name)
+
+		if t != be.GetTerminal() {
+			fail("Bundle entry [%s] doesn't match expected terminal state (exp: %v != act: %v)",
+				name, t, be.GetTerminal())
+		}
+
+		logs := exp[name]
+		for i, l := range logs {
+			if i >= len(be.GetLogs()) {
+				fail("Bundle entry [%s] missing log: %s", name, l)
+				continue
+			}
+			le := be.GetLogs()[i]
+
+			if logEntryName(le) != l {
+				fail("Bundle entry [%s] log %d doesn't match expected (exp: %s != act: %s)",
+					name, i, l, logEntryName(le))
+				continue
+			}
+		}
+		if len(be.GetLogs()) > len(logs) {
+			for _, le := range be.GetLogs()[len(logs):] {
+				fail("Bundle entry [%s] has extra log entry: %s", name, logEntryName(le))
+			}
+		}
+	}
+	for k := range entries {
+		fail("Unexpected bundle entry present: [%s]", k)
+	}
+	return strings.Join(errors, "\n")
+}
+
+func TestBundler(t *testing.T) {
+	Convey(`An empty Bundler`, t, func() {
+		b := New(Config{}).(*bundlerImpl)
+
+		Convey(`Has a size of 0 and nil GetBundles() return value.`, func() {
+			So(b.Size(), ShouldEqual, 0)
+			So(b.Empty(), ShouldBeTrue)
+			So(b.GetBundles(), ShouldBeNil)
+		})
+
+		Convey(`When adding an empty entry, still has size 0 and nil GetBundles() return value.`, func() {
+			b.Append(gen("a", false))
+			So(b.Size(), ShouldEqual, 0)
+			So(b.Empty(), ShouldBeTrue)
+			So(b.GetBundles(), ShouldBeNil)
+		})
+
+		Convey(`Bundles a terminal entry with no logs.`, func() {
+			b.Append(gen("a", true))
+
+			size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+			So(empty, ShouldBeFalse)
+
+			So(len(bundles), ShouldEqual, 1)
+			So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+			So(bundles[0], shouldHaveBundleEntries, "+a")
+		})
+
+		Convey(`Bundles an entry with 3 logs.`, func() {
+			b.Append(gen("a", false, "1", "2"))
+			b.Append(gen("a", false, "3"))
+
+			size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+			So(empty, ShouldBeFalse)
+
+			So(len(bundles), ShouldEqual, 1)
+			So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+			So(bundles[0], shouldHaveBundleEntries, "a:1:2:3")
+		})
+
+		Convey(`Bundles 2 entries with 2 logs each and one terminal entry with no logs.`, func() {
+			b.Append(gen("a", false, "1", "2"))
+			b.Append(gen("b", false, "3", "4"))
+			b.Append(gen("c", true))
+			b.Append(gen("d", false))
+
+			size, empty, bundles := b.Size(), b.Empty(), b.GetBundles()
+			So(empty, ShouldBeFalse)
+
+			So(len(bundles), ShouldEqual, 1)
+			So(size, ShouldBeGreaterThanOrEqualTo, proto.Size(bundles[0]))
+			So(bundles[0], shouldHaveBundleEntries, "a:1:2", "b:3:4", "+c")
+		})
+	})
+
+	Convey(`A Bundler with a fake Sizer`, t, func() {
+		source := "test suite"
+		dropped := []string(nil)
+		b := New(Config{
+			Threshold: 20,
+			TemplateBundle: protocol.ButlerLogBundle{
+				Source: &source,
+			},
+			Deterministic: true,
+			NewSizer: func(*protocol.ButlerLogBundle) Sizer {
+				return &fakeSizer{
+					Bundle:      8,
+					BundleEntry: 2,
+					LogEntry:    5,
+				}
+			},
+			DropCallback: func(e *protocol.ButlerLogBundle_Entry) {
+				trm := ""
+				if e.GetTerminal() {
+					trm = "+"
+				}
+				prefix := e.GetDesc().GetPrefix()
+				if len(e.GetLogs()) > 0 {
+					for _, le := range e.GetLogs() {
+						dropped = append(dropped, fmt.Sprintf("%s%s:%s", trm, prefix, le.GetLines()[0]))
+					}
+				} else {
+					dropped = append(dropped, fmt.Sprintf("%s%s", trm, prefix))
+				}
+			},
+		}).(*bundlerImpl)
+		So(b.Size(), ShouldEqual, 8)
+
+		Convey(`Adding an entry with 5 log messages outputs three bundles.`, func() {
+			b.Append(gen("a", false, "1", "2", "3", "4", "5"))
+
+			bundles := b.GetBundles()
+			So(len(bundles), ShouldEqual, 3)
+
+			Convey(`All bundles use the template bundle's fields.`, func() {
+				So(bundles[0].GetSource(), ShouldEqual, source)
+				So(bundles[1].GetSource(), ShouldEqual, source)
+				So(bundles[2].GetSource(), ShouldEqual, source)
+			})
+
+			Convey(`Have the right entries: {1,2}, {3,4}, {5}.`, func() {
+				So(bundles[0], shouldHaveBundleEntries, "a:1:2")
+				So(bundles[1], shouldHaveBundleEntries, "a:3:4")
+				So(bundles[2], shouldHaveBundleEntries, "a:5")
+			})
+		})
+
+		Convey(`Adding two entries with 2 log messages each outputs first, then second.`, func() {
+			b.Append(gen("a", false, "1", "2"))
+			b.Append(gen("b", false, "3", "4"))
+
+			bundles := b.GetBundles()
+			So(len(bundles), ShouldEqual, 2)
+			So(bundles[0], shouldHaveBundleEntries, "a:1:2")
+			So(bundles[1], shouldHaveBundleEntries, "b:3:4")
+		})
+
+		Convey(`A non-terminal entry followed by a terminal version gets output as terminal.`, func() {
+			b.Append(gen("a", false, "1"))
+			b.Append(gen("a", true, "2"))
+
+			bundles := b.GetBundles()
+			So(len(bundles), ShouldEqual, 1)
+			So(bundles[0], shouldHaveBundleEntries, "+a:1:2")
+		})
+
+		Convey(`A terminal entry followed by a non-terminal version gets output as terminal.`, func() {
+			b.Append(gen("a", true, "1"))
+			b.Append(gen("a", false, "2"))
+			b.Append(gen("a", false))
+
+			bundles := b.GetBundles()
+			So(len(bundles), ShouldEqual, 1)
+			So(bundles[0], shouldHaveBundleEntries, "+a:1:2")
+		})
+
+		Convey(`When the base bundle is above threshold, clears logs and returns nil.`, func() {
+			b.Append(gen("a", true))
+
+			So(b.Size(), ShouldEqual, 10)
+			So(b.getBundlesImpl(7), ShouldBeNil)
+
+			So(b.Size(), ShouldEqual, 8)
+			So(b.getBundlesImpl(0), ShouldBeNil)
+		})
+
+		Convey(`When the bundle entry size is above threshold, clears logs and returns nil.`, func() {
+			b.Append(gen("a", true))
+
+			So(b.Size(), ShouldEqual, 10)
+			So(b.getBundlesImpl(9), ShouldBeNil)
+
+			So(b.Size(), ShouldEqual, 8)
+			So(b.getBundlesImpl(0), ShouldBeNil)
+		})
+
+		Convey(`When the bundle has a log entry larger than threshold, it discards it.`, func() {
+			b.Append(gen("a", false, "a", "bbb", "cc", "d"))
+			b.Append(gen("b", false, "1", "2", "3"))
+
+			bundles := b.GetBundles()
+			So(len(bundles), ShouldEqual, 4)
+			So(bundles[0], shouldHaveBundleEntries, "a:a:d")
+			So(bundles[1], shouldHaveBundleEntries, "a:cc")
+			So(bundles[2], shouldHaveBundleEntries, "b:1:2")
+			So(bundles[3], shouldHaveBundleEntries, "b:3")
+			So(dropped, ShouldResemble, []string{"a:bbb"})
+
+			So(b.getBundlesImpl(0), ShouldBeNil)
+		})
+	})
+}