github.com/grafana/pyroscope@v1.18.0/pkg/api/connect/codec.go (about)

     1  package connectapi
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"connectrpc.com/connect"
     7  	"google.golang.org/protobuf/proto"
     8  )
     9  
    10  // Name is the name registered for the proto compressor.
    11  const Name = "proto"
    12  
    13  var ProtoCodec connect.Codec = vtprotoCodec{}
    14  
    15  type vtprotoCodec struct{}
    16  
    17  // growcap scales up the capacity of a slice.
    18  //
    19  // Given a slice with a current capacity of oldcap and a desired
    20  // capacity of wantcap, growcap returns a new capacity >= wantcap.
    21  //
    22  // The algorithm is mostly identical to the one used by append as of Go 1.14.
    23  func growcap(oldcap, wantcap int) (newcap int) {
    24  	if wantcap > oldcap*2 {
    25  		newcap = wantcap
    26  	} else if oldcap < 1024 {
    27  		// The Go 1.14 runtime takes this case when len(s) < 1024,
    28  		// not when cap(s) < 1024. The difference doesn't seem
    29  		// significant here.
    30  		newcap = oldcap * 2
    31  	} else {
    32  		newcap = oldcap
    33  		for 0 < newcap && newcap < wantcap {
    34  			newcap += newcap / 4
    35  		}
    36  		if newcap <= 0 {
    37  			newcap = wantcap
    38  		}
    39  	}
    40  	return newcap
    41  }
    42  
    43  type vtprotoMessage interface {
    44  	MarshalVT() ([]byte, error)
    45  	MarshalToSizedBufferVT([]byte) (int, error)
    46  	SizeVT() (n int)
    47  	UnmarshalVT([]byte) error
    48  }
    49  
    50  func (vtprotoCodec) Marshal(v any) ([]byte, error) {
    51  	switch v := v.(type) {
    52  	case vtprotoMessage:
    53  		return v.MarshalVT()
    54  	case proto.Message:
    55  		return proto.Marshal(v)
    56  	default:
    57  		return nil, fmt.Errorf("failed to marshal, message is %T, must satisfy the vtprotoMessage interface or want proto.Message", v)
    58  	}
    59  }
    60  
    61  func (vtprotoCodec) MarshalAppend(data []byte, v any) ([]byte, error) {
    62  	switch v := v.(type) {
    63  	case vtprotoMessage:
    64  		if v == nil {
    65  			return data, nil
    66  		}
    67  
    68  		n := v.SizeVT()
    69  		if cap(data) < len(data)+n {
    70  			ndata := make([]byte, len(data), growcap(cap(data), len(data)+n))
    71  			copy(ndata, data)
    72  			data = ndata
    73  		}
    74  		_, err := v.MarshalToSizedBufferVT(data[len(data) : len(data)+n])
    75  		if err != nil {
    76  			return nil, err
    77  		}
    78  		return data[:len(data)+n], nil
    79  	case proto.Message:
    80  		return proto.MarshalOptions{}.MarshalAppend(data, v)
    81  	default:
    82  		return nil, fmt.Errorf("failed to marshalAppend, message is %T, must satisfy the vtprotoMessage interface or want proto.Message", v)
    83  	}
    84  }
    85  
    86  func (vtprotoCodec) Unmarshal(data []byte, v any) error {
    87  	switch v := v.(type) {
    88  	case vtprotoMessage:
    89  		return v.UnmarshalVT(data)
    90  	case proto.Message:
    91  		return proto.Unmarshal(data, v)
    92  	default:
    93  		return fmt.Errorf("failed to unmarshal, message is %T, must satisfy the vtprotoMessage interface or want proto.Message", v)
    94  	}
    95  }
    96  
    97  func (vtprotoCodec) Name() string {
    98  	return Name
    99  }