trpc.group/trpc-go/trpc-go@v1.0.3/internal/packetbuffer/packetbuffer.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package packetbuffer implements functions for the manipulation
    15  // of byte slices.
    16  package packetbuffer
    17  
    18  import (
    19  	"fmt"
    20  	"io"
    21  	"net"
    22  
    23  	"trpc.group/trpc-go/trpc-go/internal/allocator"
    24  )
    25  
    26  // New creates a packet buffer with specific packet connection and size.
    27  func New(conn net.PacketConn, size int) *PacketBuffer {
    28  	buf, i := allocator.Malloc(size)
    29  	return &PacketBuffer{
    30  		buf:      buf,
    31  		conn:     conn,
    32  		toBeFree: i,
    33  	}
    34  }
    35  
    36  // PacketBuffer encapsulates a packet connection and implements the io.Reader interface.
    37  type PacketBuffer struct {
    38  	buf      []byte
    39  	toBeFree interface{}
    40  	conn     net.PacketConn
    41  	raddr    net.Addr
    42  	r, w     int
    43  }
    44  
    45  // Read reads data from the packet. Continuous reads cannot cross between multiple packet only if Close is called.
    46  func (pb *PacketBuffer) Read(p []byte) (n int, err error) {
    47  	if len(p) == 0 {
    48  		return 0, nil
    49  	}
    50  	if pb.w == 0 {
    51  		n, raddr, err := pb.conn.ReadFrom(pb.buf)
    52  		if err != nil {
    53  			return 0, err
    54  		}
    55  		pb.w = n
    56  		pb.raddr = raddr
    57  	}
    58  	n = copy(p, pb.buf[pb.r:pb.w])
    59  	if n == 0 {
    60  		return 0, io.EOF
    61  	}
    62  	pb.r += n
    63  	return n, nil
    64  }
    65  
    66  // Next is used to distinguish continuous logic reads. It indicates that the reading on current packet has finished.
    67  // If there remains data unconsumed, Next returns an error and discards the remaining data.
    68  func (pb *PacketBuffer) Next() error {
    69  	if pb.w == 0 {
    70  		return nil
    71  	}
    72  	var err error
    73  	if remain := pb.w - pb.r; remain != 0 {
    74  		err = fmt.Errorf("packet data is not drained, the remaining %d will be dropped", remain)
    75  	}
    76  	pb.r, pb.w = 0, 0
    77  	pb.raddr = nil
    78  	return err
    79  }
    80  
    81  // CurrentPacketAddr returns current packet's remote address.
    82  func (pb *PacketBuffer) CurrentPacketAddr() net.Addr {
    83  	return pb.raddr
    84  }
    85  
    86  // Close closes this buffer and releases resource.
    87  func (pb *PacketBuffer) Close() {
    88  	allocator.Free(pb.toBeFree)
    89  }