golib/network/server.go

package network

import (
	"net"
	"sync"
	"time"
)

type TCPListener struct {
	net.Listener
}

func (l *TCPListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	_ = conn.(*net.TCPConn).SetKeepAlivePeriod(15 * time.Second)
	_ = conn.(*net.TCPConn).SetKeepAlive(true)
	_ = conn.(*net.TCPConn).SetNoDelay(true)
	return &TCPConn{connected: true, conn: conn}, nil
}

type TCPConn struct {
	connected bool

	conn net.Conn

	// rDeadline 用于 Read 等待超时时间, 优先级高于 deadline
	rDeadline time.Time
	// wDeadline 用于 Write 等待超时时间, 优先级高于 deadline
	wDeadline time.Time
	// deadline 超时时间, 适用于 Read 和 Write, 当 rDeadline 和 wDeadline 不存在时生效
	deadline time.Time

	mu sync.Mutex
}

func (s *TCPConn) Read(b []byte) (n int, err error) {
	if !s.connected {
		return 0, ErrClosed
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err = setReadDeadline(s.conn, s.rDeadline, s.deadline); err != nil {
		return 0, err
	}
	if cap(b) == 0 {
		b = defaultPool.Get().([]byte)
		defaultPool.Put(b)
	}
	return s.conn.Read(b)
}

func (s *TCPConn) Write(b []byte) (n int, err error) {
	if !s.connected {
		return 0, ErrClosed
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err = setWriteDeadline(s.conn, s.wDeadline, s.deadline); err != nil {
		return 0, err
	}
	return s.conn.Write(b)
}

func (s *TCPConn) Close() error {
	s.mu.Lock()
	s.connected = false
	err := s.conn.Close()
	s.mu.Unlock()
	return err
}

func (s *TCPConn) LocalAddr() net.Addr {
	return s.conn.LocalAddr()
}

func (s *TCPConn) RemoteAddr() net.Addr {
	return s.conn.RemoteAddr()
}

func (s *TCPConn) SetDeadline(t time.Time) error {
	s.deadline = t
	return nil
}

func (s *TCPConn) SetReadDeadline(t time.Time) error {
	s.rDeadline = t
	return nil
}

func (s *TCPConn) SetWriteDeadline(t time.Time) error {
	s.wDeadline = t
	return nil
}