github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/netutil/conn.go (about)

     1  package netutil
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net"
     7  	"time"
     8  )
     9  
    10  func ReadBytesWithLen(c net.Conn, bytesLen int64, timeout time.Duration) ([]byte, error) {
    11  	c.SetReadDeadline(time.Now().Add(timeout))
    12  	readBytes := make([]byte, 0, bytesLen)
    13  	for bytesLen > 0 {
    14  		buf := make([]byte, bytesLen, bytesLen)
    15  		n, err := c.Read(buf)
    16  		if err != nil {
    17  			if err == io.EOF {
    18  				return readBytes, nil
    19  			}
    20  			return nil, err
    21  		}
    22  
    23  		bytesLen -= int64(n)
    24  		if n == 0 && bytesLen != 0 {
    25  			return nil, fmt.Errorf("cannot read the %d length of bytes", bytesLen)
    26  		}
    27  		readBytes = append(readBytes, buf[0:n]...)
    28  	}
    29  	return readBytes, nil
    30  }
    31  
    32  func ConnectWithRetry(network, addr string, timeout time.Duration, retryNum int) (c net.Conn, err error) {
    33  	for i := 0; i < retryNum; i++ {
    34  		c, err = net.DialTimeout(
    35  			network,
    36  			addr,
    37  			timeout,
    38  		)
    39  		if err == nil {
    40  			return
    41  		}
    42  	}
    43  	return
    44  }
    45  
    46  func WriteWithTimeout(c net.Conn, requestBytes []byte, timeout time.Duration) error {
    47  	c.SetWriteDeadline(time.Now().Add(time.Duration(timeout)))
    48  	writeLen := len(requestBytes)
    49  	for writeLen > 0 {
    50  		n, err := c.Write(requestBytes)
    51  		if err != nil {
    52  			return err
    53  		}
    54  
    55  		if n <= 0 {
    56  			err = fmt.Errorf("write 0 bytes to connection for pab")
    57  			return err
    58  		}
    59  		writeLen -= n
    60  	}
    61  	return nil
    62  }