github.com/weedge/lib@v0.0.0-20230424045628-a36dcc1d90e4/poller/codec_raw.go (about)

     1  package poller
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  type headerLenDecoder struct {
    10  	// headerlen
    11  	// Length of the TCP packet header, which is used to describe the byte length of the packet
    12  	headerLen int
    13  }
    14  
    15  // NewHeaderLenDecoder Creates a decoder based on header length
    16  // headerLen Indicates the header of a TCP packet, which describes the byte length of the packet
    17  // readMaxLen Specifies the maximum length of the client packet read. The length of the packet sent by the client cannot exceed this length
    18  func NewHeaderLenDecoder(headerLen int) Decoder {
    19  	if headerLen <= 0 {
    20  		panic("headerLen or readMaxLen must must greater than 0")
    21  	}
    22  
    23  	return &headerLenDecoder{
    24  		headerLen: headerLen,
    25  	}
    26  }
    27  
    28  // Decode
    29  func (d *headerLenDecoder) DecodeBuffer(buffer *Buffer) (value []byte, err error) {
    30  	value = []byte{}
    31  	header, err := buffer.Seek(d.headerLen)
    32  	if err != nil {
    33  		return
    34  	}
    35  
    36  	valueLen := int(binary.BigEndian.Uint16(header))
    37  	value, err = buffer.Read(d.headerLen, valueLen)
    38  	if err != nil {
    39  		return
    40  	}
    41  
    42  	return
    43  }
    44  
    45  func (d *headerLenDecoder) Decode(buffer *Buffer) (value []byte, err error) {
    46  	value, err = d.DecodeBuffer(buffer)
    47  	if err == ErrBufferNotEnough {
    48  		return []byte{}, nil
    49  	}
    50  	if err != nil {
    51  		return
    52  	}
    53  
    54  	return
    55  }
    56  
    57  type headerLenEncoder struct {
    58  	// headerLen
    59  	// Length of the TCP packet header, which is used to describe the byte length of the packet
    60  	headerLen int
    61  	// writeBufferLen
    62  	// The recommended length of packets sent by the server to the client. Writebufferlen int takes advantage of pool optimization when packets are sent less than this value
    63  	writeBufferLen  int
    64  	writeBufferPool *sync.Pool
    65  }
    66  
    67  // NewHeaderLenEncoder
    68  // Creates an encoder based on header length
    69  // headerLen Indicates the header of a TCP packet, which describes the byte length of the packet
    70  // writeBufferLen Indicates the recommended length of a packet sent by the server to the client. When a packet is sent less than this value, it takes advantage of the pool optimization
    71  func NewHeaderLenEncoder(headerLen, writeBufferLen int) *headerLenEncoder {
    72  	if headerLen <= 0 || writeBufferLen <= 0 {
    73  		panic("headerLen or writeBufferLen must must greater than 0")
    74  	}
    75  
    76  	return &headerLenEncoder{
    77  		headerLen:      headerLen,
    78  		writeBufferLen: writeBufferLen,
    79  		writeBufferPool: &sync.Pool{
    80  			New: func() interface{} {
    81  				b := make([]byte, writeBufferLen)
    82  				return b
    83  			},
    84  		},
    85  	}
    86  }
    87  
    88  // EncodeToWriter
    89  // Encodes data and writes it to Writer
    90  func (e headerLenEncoder) EncodeToWriter(w io.Writer, bytes []byte) error {
    91  	l := len(bytes)
    92  	var buffer []byte
    93  	if l <= e.writeBufferLen-e.headerLen {
    94  		obj := e.writeBufferPool.Get()
    95  		defer e.writeBufferPool.Put(obj)
    96  		buffer = obj.([]byte)[0 : l+e.headerLen]
    97  	} else {
    98  		buffer = make([]byte, l+e.headerLen)
    99  	}
   100  
   101  	// Write the message length to buffer
   102  	binary.BigEndian.PutUint16(buffer[0:e.headerLen], uint16(l))
   103  	// Writes the message content to the buffer
   104  	copy(buffer[e.headerLen:], bytes)
   105  
   106  	_, err := w.Write(buffer)
   107  	return err
   108  }