github.com/grafana/pyroscope@v1.18.0/pkg/metastore/fsm/log_entry.go (about)

     1  package fsm
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"reflect"
     7  
     8  	"google.golang.org/protobuf/proto"
     9  )
    10  
    11  var ErrInvalidCommand = fmt.Errorf("invalid command format; expected at least 4 bytes")
    12  
    13  type RaftLogEntryType uint32
    14  
    15  type RaftLogEntry struct {
    16  	Type RaftLogEntryType
    17  	Data []byte
    18  }
    19  
    20  type Response struct {
    21  	Data proto.Message
    22  	Err  error
    23  }
    24  
    25  func MarshalEntry(t RaftLogEntryType, payload proto.Message) ([]byte, error) {
    26  	b, err := marshal(payload)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  	binary.BigEndian.PutUint32(b, uint32(t))
    31  	return b, nil
    32  }
    33  
    34  func (c *RaftLogEntry) UnmarshalBinary(b []byte) error {
    35  	if len(b) < 4 {
    36  		return ErrInvalidCommand
    37  	}
    38  	c.Type = RaftLogEntryType(binary.BigEndian.Uint32(b))
    39  	c.Data = b[4:]
    40  	return nil
    41  }
    42  
    43  type vtMarshaller interface {
    44  	MarshalToSizedBufferVT([]byte) (int, error)
    45  	SizeVT() int
    46  }
    47  
    48  // marshal MUST allocate a buffer of size covering
    49  // the whole message, including the entry header.
    50  func marshal(v proto.Message) ([]byte, error) {
    51  	if m, ok := any(v).(vtMarshaller); ok {
    52  		size := m.SizeVT()
    53  		buf := make([]byte, size+4)
    54  		_, err := m.MarshalToSizedBufferVT(buf[4:])
    55  		return buf, err
    56  	}
    57  	raw, err := proto.Marshal(v)
    58  	if err != nil {
    59  		return raw, err
    60  	}
    61  	buf := make([]byte, 4+len(raw))
    62  	copy(buf[4:], raw)
    63  	return buf, err
    64  }
    65  
    66  type vtUnmarshaler interface {
    67  	UnmarshalVT([]byte) error
    68  }
    69  
    70  func newProto[T proto.Message]() T {
    71  	var msg T
    72  	msgType := reflect.TypeOf(msg).Elem()
    73  	return reflect.New(msgType).Interface().(T)
    74  }
    75  
    76  func unmarshal[T proto.Message](b []byte) (v T, err error) {
    77  	v = newProto[T]()
    78  	if vt, ok := any(v).(vtUnmarshaler); ok {
    79  		err = vt.UnmarshalVT(b)
    80  	} else {
    81  		err = proto.Unmarshal(b, v)
    82  	}
    83  	return v, err
    84  }