logdog: Add bundler library.

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

+// 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
+
+        undoSize   int64
+        undoPath   *types.StreamPath
+        undoBeSize 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 {
+                if v > 0 {
+                        size += int64(varintLength(uint64(v)))
+                }
+        }
+        return size
+}
+
+func (b *fastSizer) Append(be *protocol.ButlerLogBundle_Entry, e *protocol.LogEntry) {
+        path := protoutil.DescriptorPath(be.GetDesc())
+        size := int64(0)
+        if !b.hasStream(path) {
+                // This is the first time we've seen this ButlerLogBundle_Entry. Add its
+                // static cost.
+                size += int64(bundleEntryTagSize + proto.Size(be))
+        }
+
+        if e != nil {
+                ps := proto.Size(e)
+                size += int64(logEntryTagSize + varintLength(uint64(ps)) + ps)
+        }
+
+        b.size += size
+        b.beSize[path] += size
+
+        b.undoSize = size
+        b.undoPath = &path
+        b.undoBeSize = size
+}
+
+func (b *fastSizer) Undo() {
+        b.size -= b.undoSize
+        if b.undoPath != nil {
+                b.beSize[*b.undoPath] -= b.undoBeSize
+        }
+
+        b.undoSize = 0
+        b.undoPath = nil
+        b.undoBeSize = 0
+}
+
+func (b *fastSizer) hasStream(path types.StreamPath) bool {
+        _, ok := b.beSize[path]
+        return ok
+}
+
+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
+}