github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/xnet/timeout.go (about)

     1  package xnet
     2  
     3  import (
     4  	"net"
     5  	"time"
     6  )
     7  
     8  type timeoutListener struct {
     9  	*net.TCPListener
    10  	timeout time.Duration
    11  }
    12  
    13  func (l timeoutListener) Accept() (net.Conn, error) {
    14  	conn, err := l.TCPListener.Accept()
    15  	if err != nil {
    16  		return nil, err
    17  	}
    18  	SetTimeout(conn, l.timeout)
    19  	return conn, nil
    20  }
    21  
    22  func TimeoutListener(l net.Listener, timeout time.Duration) net.Listener {
    23  	if tl, ok := l.(*net.TCPListener); ok {
    24  		return timeoutListener{TCPListener: tl, timeout: timeout}
    25  	}
    26  	return l
    27  }
    28  
    29  func SetTimeout(c net.Conn, timeout time.Duration) {
    30  	tc, ok := c.(*net.TCPConn)
    31  	if !ok {
    32  		return
    33  	}
    34  	deadline := time.Now().Add(timeout)
    35  	tc.SetDeadline(deadline)
    36  	tc.SetReadDeadline(deadline)
    37  	tc.SetWriteDeadline(deadline)
    38  }