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

     1  package xnet
     2  
     3  import (
     4  	"net"
     5  	"time"
     6  )
     7  
     8  // keepaliveListener wraps *net.TCPListener.
     9  type keepaliveListener struct {
    10  	*net.TCPListener
    11  	timeout time.Duration
    12  }
    13  
    14  // Accept returns a TCP keep-alive enabled connection.
    15  func (l keepaliveListener) Accept() (c net.Conn, err error) {
    16  	conn, err := l.AcceptTCP()
    17  	if err != nil {
    18  		return
    19  	}
    20  	SetKeepAlive(conn, l.timeout)
    21  	return conn, nil
    22  }
    23  
    24  // KeepAliveListener returns listener if l is a
    25  // *net.TCPListener.  Otherwise, l is returned without change.
    26  func KeepAliveListener(l net.Listener, keepalive time.Duration) net.Listener {
    27  	if tl, ok := l.(*net.TCPListener); ok {
    28  		return keepaliveListener{TCPListener: tl, timeout: keepalive}
    29  	}
    30  	return l
    31  }
    32  
    33  // SetKeepAlive enables TCP keep-alive if c is a *net.TCPConn.
    34  func SetKeepAlive(c net.Conn, keepalive time.Duration) {
    35  	tc, ok := c.(*net.TCPConn)
    36  	if !ok {
    37  		return
    38  	}
    39  	tc.SetKeepAlive(true)
    40  	tc.SetKeepAlivePeriod(keepalive)
    41  }