logdog: Add bundler library.

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

Index: client/internal/logdog/butler/bundler/sizer_fast.go
diff --git a/client/internal/logdog/butler/bundler/sizer_fast.go b/client/internal/logdog/butler/bundler/sizer_fast.go
new file mode 100644
index 0000000000000000000000000000000000000000..3709e75a04b54d9105ace18f059e1b1503b016a0
--- /dev/null
+++ b/client/internal/logdog/butler/bundler/sizer_fast.go
@@ -0,0 +1,153 @@
+// 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 (
+	"errors"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+
+	"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"
+)
+
+var (
+	bundleEntryTagSize int
+	logEntryTagSize    int
+
+	errMalformedProtobufField = errors.New("malformed protobuf field")
+)
+
+func init() {
+	bundleEntryTagSize = mustCalculateTagSize(&protocol.ButlerLogBundle{}, "Entries")
+	logEntryTagSize = mustCalculateTagSize(&protocol.ButlerLogBundle_Entry{}, "Logs")
+}
+
+type fastSizer struct {
+	size int64
+
+	// As we add LogEntry to a ButlerLogBundle_Entry, the bundle entry's size (and
+	// therefore it's protobuf size prefix) will grow. We account for the growth
+	// by tracking the size of each ButlerLogBundle_Entry and updating it when
+	// we add a LogEntry to it. This is stored independently from the cumulative
+	// size and factored in when Size() is calculated.
+	beSize map[types.StreamPath]int64
+}
+
+// NewFastSizer is a Sizer that is optimized for LogDog protobufs.
+//
+// In exchange for rapid size calculation, it performs worst-case estimates on
+// the unknown protocol overheads, leading to potential size overestimation.
+func NewFastSizer(b *protocol.ButlerLogBundle) Sizer {
+	return &fastSizer{
+		size:   int64(proto.Size(b)),
+		beSize: map[types.StreamPath]int64{},
+	}
+}
+
+func (b *fastSizer) Size() int64 {
+	size := b.size
+	for _, v := range b.beSize {
+		size += int64(varintLength(uint64(v)))
+	}
+	return size
+}
+
+func (b *fastSizer) AppendBundleEntry(e *protocol.ButlerLogBundle_Entry) {
+	ps := proto.Size(e)
+	s := bundleEntryTagSize + ps
+	b.size += int64(s)
+	b.addBundleEntrySize(e, s)
+}
+
+func (b *fastSizer) AppendLogEntry(be *protocol.ButlerLogBundle_Entry, e *protocol.LogEntry) {
+	ps := proto.Size(e)
+	s := logEntryTagSize + varintLength(uint64(ps)) + ps
+	b.size += int64(s)
+	b.addBundleEntrySize(be, s)
+}
+
+func (b *fastSizer) addBundleEntrySize(e *protocol.ButlerLogBundle_Entry, count int) {
+	path := protoutil.DescriptorPath(e.GetDesc())
+	b.beSize[path] += int64(count)
+}
+
+func mustCalculateTagSize(i interface{}, field string) int {
+	value, err := calculateTagSize(i, field)
+	if err != nil {
+		panic(err)
+	}
+	return value
+}
+
+func calculateTagSize(i interface{}, field string) (int, error) {
+	v := reflect.TypeOf(i)
+	if v.Kind() == reflect.Ptr {
+		v = v.Elem()
+	}
+	if v.Kind() != reflect.Struct {
+		return 0, fmt.Errorf("sizer: %s is not a struct", v)
+	}
+
+	f, ok := v.FieldByName(field)
+	if !ok {
+		return 0, fmt.Errorf("sizer: could not find field %s.%s", v, field)
+	}
+
+	tag, err := protobufTag(f)
+	if err != nil {
+		return 0, fmt.Errorf("sizer: field %s.%s has no protobuf tag: %s", v, field, err)
+	}
+
+	// Protobuf encodes the tag and wire type in the same varint. It does this
+	// by allocating three bits for wire type at the base of the tag.
+	//
+	// https://developers.google.com/protocol-buffers/docs/encoding#structure
+	return varintLength(uint64(tag) << 3), nil
+}
+
+func varintLength(val uint64) int {
+	switch {
+	case val < 0x80:
+		return 1
+	case val < 0x4000:
+		return 2
+	case val < 0x200000:
+		return 3
+	case val < 0x10000000:
+		return 4
+	case val < 0x800000000:
+		return 5
+	case val < 0x40000000000:
+		return 6
+	case val < 0x2000000000000:
+		return 7
+	case val < 0x100000000000000:
+		return 8
+	case val < 0x8000000000000000:
+		return 9
+	default:
+		// Maximum uvarint size.
+		return 10
+	}
+}
+
+func protobufTag(f reflect.StructField) (int, error) {
+	// If this field doesn't have a "protobuf" tag, ignore it.
+	value := f.Tag.Get("protobuf")
+	parts := strings.Split(value, ",")
+	if len(parts) < 2 {
+		return 0, errMalformedProtobufField
+	}
+	tag, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return 0, errMalformedProtobufField
+	}
+	return tag, nil
+}