github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/protocol/message/mux/serialization.go (about)

     1  package mux
     2  
     3  import (
     4  	"encoding/binary"
     5  	"github.com/nyan233/littlerpc/core/container"
     6  	. "github.com/nyan233/littlerpc/core/protocol"
     7  	message2 "github.com/nyan233/littlerpc/core/protocol/message"
     8  )
     9  
    10  // Marshal 与message.Marshal不同, 该函数没有任何副作用
    11  func Marshal(msg *Block, payloads *container.Slice[byte]) {
    12  	start := payloads.Len()
    13  	payloads.Append(make([]byte, BlockBaseLen))
    14  	(*payloads)[start+0] = msg.Flags
    15  	binary.BigEndian.PutUint32((*payloads)[start+1:start+5], msg.StreamId)
    16  	binary.BigEndian.PutUint64((*payloads)[start+5:start+13], msg.MsgId)
    17  	binary.BigEndian.PutUint16((*payloads)[start+13:start+15], msg.PayloadLength)
    18  	payloads.Append(msg.Payloads)
    19  	return
    20  }
    21  
    22  func Unmarshal(data container.Slice[byte], msg *Block) error {
    23  	if data.Len() < BlockBaseLen {
    24  		return ErrBadMessage
    25  	}
    26  	msg.Flags = data[0]
    27  	data = data[1:]
    28  	msg.StreamId = binary.BigEndian.Uint32(data[:4])
    29  	data = data[4:]
    30  	msg.MsgId = binary.BigEndian.Uint64(data[:8])
    31  	data = data[8:]
    32  	msg.PayloadLength = binary.BigEndian.Uint16(data[:2])
    33  	msg.Payloads.Reset()
    34  	msg.Payloads = data[2:]
    35  	return nil
    36  }
    37  
    38  // MarshalIteratorFromMessage buf1在返回时可以回收, buf2需要迭代器完成工作时才可回收
    39  // base Block中需要有除Payloads&PayloadLength之外的所有信息
    40  func MarshalIteratorFromMessage(msg *message2.Message, noMuxBuf, muxBuf *container.Slice[byte], base Block) (*container.Iterator[[]byte], error) {
    41  	noMuxBuf.Reset()
    42  	muxBuf.Reset()
    43  	var nBlock int
    44  	if err := message2.Marshal(msg, noMuxBuf); err != nil {
    45  		return nil, err
    46  	}
    47  	for noMuxBuf.Len() > 0 {
    48  		nBlock++
    49  		var copyLength int
    50  		if noMuxBuf.Len() > MaxPayloadSizeOnMux {
    51  			copyLength = MaxPayloadSizeOnMux
    52  		} else {
    53  			copyLength = noMuxBuf.Len()
    54  		}
    55  		base.SetPayloads((*noMuxBuf)[:copyLength])
    56  		*noMuxBuf = (*noMuxBuf)[copyLength:]
    57  		Marshal(&base, muxBuf)
    58  	}
    59  	iter := container.NewIterator[[]byte](nBlock, true, func(current int) []byte {
    60  		start := current * MaxBlockSize
    61  		payloadLength := binary.BigEndian.Uint16((*muxBuf)[start+13 : start+15])
    62  		return (*muxBuf)[start : start+15+int(payloadLength)]
    63  	}, func() {
    64  		return
    65  	})
    66  	return iter, nil
    67  }