myth/tcp/io.go

package tcp

import (
	"bufio"
	"net"
	"time"
)

func errTimeout(err error) bool {
	if opErr, ok := err.(*net.OpError); ok {
		if opErr.Temporary() || opErr.Timeout() {
			return true
		}
	}
	return false
}

// Write tcp write
func Write(conn *net.TCPConn, timeout time.Duration, data []byte) (n int, isTimeout bool, err error) {
	conn.SetWriteDeadline(time.Now().Add(timeout))
	n, err = conn.Write(data)
	if err != nil {
		isTimeout = errTimeout(err)
	}
	return
}

// Read tcp reads data into p.
// It returns the number of bytes read into p.
// The bytes are taken from at most one Read on the underlying Reader,
// hence n may be less than len(p).
// To read exactly len(p) bytes, use io.ReadFull(b, p).
// At EOF, the count will be zero and err will be io.EOF.
func Read(conn *net.TCPConn, timeout time.Duration, p []byte) (n int, isTimeout bool, err error) {
	buffer := bufio.NewReader(conn)
	conn.SetReadDeadline(time.Now().Add(timeout))

	n, err = buffer.Read(p)
	if err != nil {
		isTimeout = errTimeout(err)
	}
	return
}

// ReadBytes tcp reads until the first occurrence of delim in the input,
// returning a slice containing the data up to and including the delimiter.
// If ReadBytes encounters an error before finding a delimiter,
// it returns the data read before the error and the error itself (often io.EOF).
// ReadBytes returns err != nil if and only if the returned data does not end in
// delim.
// For simple uses, a Scanner may be more convenient.
func ReadBytes(conn *net.TCPConn, timeout time.Duration, delim byte) (data []byte, isTimeout bool, err error) {
	buffer := bufio.NewReader(conn)
	conn.SetReadDeadline(time.Now().Add(timeout))

	data, err = buffer.ReadBytes(delim)
	if err != nil {
		isTimeout = errTimeout(err)
	}
	return
}

// Optimize TCP connect optimize
func Optimize(conn *net.TCPConn) (err error) {
	err = conn.SetNoDelay(true)
	if err != nil {
		return
	}
	err = conn.SetKeepAlive(true)
	if err != nil {
		return
	}
	err = conn.SetKeepAlivePeriod(3 * time.Minute)
	return
}