MrBlue

package util

import (
    "fmt"
    "reflect"

    "github.com/golang/protobuf/proto"
)

//------------------------------------------------------
// 分析proto结构，计算每个字段序列化后大小
//------------------------------------------------------

// fieldInfo 字段信息
type fieldInfo struct {
    Name  string // 字段名字
    Type  string // 字段类型
    Count int    // 个数(map和slice可能是多个)
    Size  int    // pb序列化成后的大小
}

// ProtoStructInfo // pb结构体信息
type ProtoStructInfo struct {
    Tag        string // 自定义标记
    Total      int    // 序列化后总大小
    Accumulate int    // 每个字段序列化后累加大小
    Fields     fields // 每个字段信息
}

type fields []fieldInfo

// sortFunc quick sort
func sortFunc(f fields, i, j int) {
    if i >= j {
        return
    }
    b, e := i, j
    key := f[i]
    for i < j {
        for i < j && key.Size >= f[j].Size {
            j--
        }
        if i < j {
            f[i] = f[j]
        }
        for i < j && key.Size <= f[i].Size {
            i++
        }
        if i < j {
            f[j] = f[i]
        }
    }
    f[i] = key
    sortFunc(f, b, i-1)
    sortFunc(f, i+1, e)
}

func (f fields) sort() {
    if 0 == len(f) {
        return
    }

    sortFunc(f, 0, len(f)-1)
}

// String 打印字符串
func (f ProtoStructInfo) String() string {
    str := ""
    if f.Tag != "" {
        str = fmt.Sprintf("%s\n", f.Tag)
    }
    str = str + fmt.Sprintf("total\t%8d\nsum\t\t%8d\n", f.Total, f.Accumulate)

    n := 0
    for k, v := range f.Fields {
        //str += fmt.Sprintf("%-16s\t%-16s\tcount:%10d\tsize:%10d\n", v.Name, v.Type, v.Count, v.Size)
        n += v.Size
        str += fmt.Sprintf("%2d %-16s\t %-28s\tcount\t%4d\tsize\t%6d\tflat%%\t%.2f%%\tsum\t%6d\tsum%%\t%.2f%%\n", k, v.Name, v.Type, v.Count, v.Size, 100*float32(v.Size)/float32(f.Accumulate), n, 100*float32(n)/float32(f.Accumulate))
    }
    return str
}

// AnalyseProto 分析proto结构
func AnalyseProto(ptr proto.Message) ProtoStructInfo {
    ret := ProtoStructInfo{
        Fields: make([]fieldInfo, 0),
    }

    // 计算总大小
    if b, err := proto.Marshal(ptr); nil != err {
        panic(err)
    } else {
        ret.Total = len(b)
    }

    v := reflect.Indirect(reflect.ValueOf(ptr))
    t := v.Type()

    for i := 0; i < v.NumField(); i++ {
        f := v.Field(i)
        tt := t.Field(i)
        // 只累计这几个类型，其他类型没法单独计算大小
        if f.Kind() != reflect.Ptr && f.Kind() != reflect.Map && f.Kind() != reflect.Slice {
            continue
        }
        if f.IsNil() {
            continue
        }

        n := 0     // 字段大小
        count := 0 // 字段个数
        if f.Kind() == reflect.Ptr {
            if m, ok := f.Interface().(proto.Message); !ok {
                continue
            } else {
                b, err := proto.Marshal(m)
                if nil != err {
                    panic(err)
                }
                n = len(b)
                count++
            }
        } else if f.Kind() == reflect.Map {
            for it := f.MapRange(); it.Next(); {
                count++
                if m, ok := it.Value().Interface().(proto.Message); !ok {
                    continue
                } else {
                    b, err := proto.Marshal(m)
                    if nil != err {
                        panic(err)
                    }
                    n += len(b)

                }
            }
        } else if f.Kind() == reflect.Slice {
            for i := 0; i < f.Len(); i++ {
                ff := f.Index(i)
                count++
                if m, ok := ff.Interface().(proto.Message); !ok {
                    break
                } else {
                    b, err := proto.Marshal(m)
                    if nil != err {
                        panic(err)
                    }
                    n += len(b)

                }
            }
        }

        if 0 == n && 0 == count {
            continue
        }
        fi := fieldInfo{
            Name:  tt.Name,
            Type:  f.Type().String(),
            Count: count,
            Size:  n,
        }
        ret.Fields = append(ret.Fields, fi)
        ret.Accumulate += n
    }
    ret.Fields.sort()
    return ret
}

fmt.Println(util.AnalyseProto(&pb.User{Name:1,ID:2}))