使用io.Copy简化Golang TCP代理实现的思考

以前尝试给Easegress中添加四层代理功能时，曾面临双向连接错误控制的挑战。
当时的想法是每条TCP连接分别使用两个goroutine进行读写，通过buf实现流量拷贝，并在读写出错时判断应仅关闭读端、写端，还是整个连接。
此外还需处理两端TCP连接因读写能力不匹配而引发的流量控制问题。

1 初版实现

package tcpproxy

import (
    "io"
    "net"
    "runtime/debug"
    "sync"
    "sync/atomic"
    "time"

    "github.com/megaease/easegress/pkg/logger"
    "github.com/megaease/easegress/pkg/util/fasttime"
    "github.com/megaease/easegress/pkg/util/iobufferpool"
    "github.com/megaease/easegress/pkg/util/timerpool"
)

const writeBufSize = 8

var tcpBufferPool = sync.Pool{
    New: func() interface{} {
        buf := make([]byte, iobufferpool.DefaultBufferReadCapacity)
        return buf
    },
}

// Connection wrap tcp connection
type Connection struct {
    closed     uint32
    rawConn    net.Conn
    localAddr  net.Addr
    remoteAddr net.Addr

    readBuffer      []byte
    writeBuffers    net.Buffers
    ioBuffers       []*iobufferpool.StreamBuffer
    writeBufferChan chan *iobufferpool.StreamBuffer

    mu               sync.Mutex
    connStopChan     chan struct{} // use for connection close
    listenerStopChan chan struct{} // use for listener close

    lastReadDeadlineTime  time.Time
    lastWriteDeadlineTime time.Time

    onRead  func(buffer *iobufferpool.StreamBuffer) // execute read filters
    onClose func(event ConnectionEvent)
}

// NewClientConn wrap connection create from client
func NewClientConn(conn net.Conn, listenerStopChan chan struct{}) *Connection {
    return &Connection{
        rawConn:          conn,
        localAddr:        conn.LocalAddr(),
        remoteAddr:       conn.RemoteAddr(),
        listenerStopChan: listenerStopChan,

        mu:              sync.Mutex{},
        connStopChan:    make(chan struct{}),
        writeBufferChan: make(chan *iobufferpool.StreamBuffer, writeBufSize),
    }
}

// SetOnRead set connection read handle
func (c *Connection) SetOnRead(onRead func(buffer *iobufferpool.StreamBuffer)) {
    c.onRead = onRead
}

// SetOnClose set close callback
func (c *Connection) SetOnClose(onclose func(event ConnectionEvent)) {
    c.onClose = onclose
}

// Start running connection read/write loop
func (c *Connection) Start() {
    fnRecover := func() {
        if r := recover(); r != nil {
            logger.Errorf("tcp read/write loop panic: %v\n%s\n", r, string(debug.Stack()))
            c.Close(NoFlush, LocalClose)
        }
    }

    go func() {
        defer fnRecover()
        c.startReadLoop()
    }()

    go func() {
        defer fnRecover()
        c.startWriteLoop()
    }()
}

// Write receive other connection data
func (c *Connection) Write(buf *iobufferpool.StreamBuffer) (err error) {
    defer func() {
        if r := recover(); r != nil {
            logger.Errorf("tcp connection has closed, local addr: %s, remote addr: %s, err: %+v",
                c.localAddr.String(), c.remoteAddr.String(), r)
            err = ErrConnectionHasClosed
        }
    }()

    select {
    case c.writeBufferChan <- buf:
        return
    default:
    }

    // try to send data again in 60 seconds
    t := timerpool.Get(60 * time.Second)
    select {
    case c.writeBufferChan <- buf:
    case <-t.C:
        buf.Release()
        err = ErrWriteBufferChanTimeout
    }
    timerpool.Put(t)
    return
}

func (c *Connection) startReadLoop() {
    defer func() {
        if c.readBuffer != nil {
            tcpBufferPool.Put(c.readBuffer[:iobufferpool.DefaultBufferReadCapacity])
        }
    }()

    for {
        select {
        case <-c.connStopChan:
            return
        case <-c.listenerStopChan:
            logger.Debugf("connection close due to listener stopped, local addr: %s, remote addr: %s",
                c.localAddr.String(), c.remoteAddr.String())
            c.Close(NoFlush, LocalClose)
            return
        default:
        }

        n, err := c.doReadIO()
        if n > 0 {
            c.onRead(iobufferpool.NewStreamBuffer(c.readBuffer[:n]))
        }

        if err == nil {
            continue
        }

        if te, ok := err.(net.Error); ok && te.Timeout() {
            select {
            case <-c.connStopChan:
                logger.Debugf("connection has closed, exit read loop, local addr: %s, remote addr: %s",
                    c.localAddr.String(), c.remoteAddr.String())
                return
            default:
            }
            continue // ignore timeout error, continue read data
        }

        if err == io.EOF {
            logger.Debugf("remote close connection, local addr: %s, remote addr: %s, err: %s",
                c.localAddr.String(), c.remoteAddr.String(), err.Error())
            c.Close(NoFlush, RemoteClose)
        } else {
            logger.Errorf("error on read, local addr: %s, remote addr: %s, err: %s",
                c.localAddr.String(), c.remoteAddr.String(), err.Error())
            c.Close(NoFlush, OnReadErrClose)
        }
        return
    }
}

func (c *Connection) startWriteLoop() {
    var err error
    for {
        select {
        case <-c.connStopChan:
            logger.Debugf("connection exit write loop, local addr: %s, remote addr: %s",
                c.localAddr.String(), c.remoteAddr.String())
            return
        case buf, ok := <-c.writeBufferChan:
            if !ok {
                return
            }
            c.appendBuffer(buf)
        NoMoreData:
            // Keep reading until writeBufferChan is empty
            // writeBufferChan may be full when writeLoop call doWrite
            for i := 0; i < writeBufSize-1; i++ {
                select {
                case buf, ok := <-c.writeBufferChan:
                    if !ok {
                        return
                    }
                    c.appendBuffer(buf)
                default:
                    break NoMoreData
                }
            }
        }

        if _, err = c.doWrite(); err == nil {
            continue
        }

        if te, ok := err.(net.Error); ok && te.Timeout() {
            select {
            case <-c.connStopChan:
                logger.Debugf("connection has closed, exit write loop, local addr: %s, remote addr: %s",
                    c.localAddr.String(), c.remoteAddr.String())
                return
            default:
            }

            c.Close(NoFlush, OnWriteTimeout)
            return
        }

        if err == iobufferpool.ErrEOF {
            logger.Debugf("finish write with eof, local addr: %s, remote addr: %s",
                c.localAddr.String(), c.remoteAddr.String())
            c.Close(NoFlush, LocalClose)
        } else {
            // remote call CloseRead, so just exit write loop, wait read loop exit
            logger.Errorf("error on write, local addr: %s, remote addr: %s, err: %+v",
                c.localAddr.String(), c.remoteAddr.String(), err)
        }
        return
    }
}

func (c *Connection) appendBuffer(buf *iobufferpool.StreamBuffer) {
    if buf == nil {
        return
    }
    c.ioBuffers = append(c.ioBuffers, buf)
    c.writeBuffers = append(c.writeBuffers, buf.Bytes())
}

// Close connection close function
func (c *Connection) Close(ccType CloseType, event ConnectionEvent) {
    defer func() {
        if r := recover(); r != nil {
            logger.Errorf("connection close panic, err: %+v\n%s", r, string(debug.Stack()))
        }
    }()

    if ccType == FlushWrite {
        _ = c.Write(iobufferpool.NewEOFStreamBuffer())
        return
    }

    if !atomic.CompareAndSwapUint32(&c.closed, 0, 1) {
        // connection has already closed, so there is no need to execute below code
        return
    }

    // close tcp rawConn read first
    logger.Debugf("enter connection close func(%s), local addr: %s, remote addr: %s",
        event, c.localAddr.String(), c.remoteAddr.String())

    close(c.connStopChan)
    _ = c.rawConn.SetDeadline(time.Now()) // notify break read/write loop

    c.onClose(event)
    _ = c.rawConn.Close()
}

func (c *Connection) doReadIO() (bufLen int, err error) {
    if c.readBuffer == nil {
        c.readBuffer = tcpBufferPool.Get().([]byte)
    }

    // add read deadline setting optimization?
    // https://github.com/golang/go/issues/15133
    curr := fasttime.Now().Add(15 * time.Second)
    // there is no need to set readDeadline in too short time duration
    if diff := curr.Sub(c.lastReadDeadlineTime).Milliseconds(); diff > 0 {
        _ = c.rawConn.SetReadDeadline(curr)
        c.lastReadDeadlineTime = curr
    }
    return c.rawConn.(io.Reader).Read(c.readBuffer)
}

func (c *Connection) doWrite() (int64, error) {
    curr := fasttime.Now().Add(15 * time.Second)
    // there is no need to set writeDeadline in too short time duration
    if diff := curr.Sub(c.lastWriteDeadlineTime).Milliseconds(); diff > 0 {
        _ = c.rawConn.SetWriteDeadline(curr)
        c.lastWriteDeadlineTime = curr
    }
    return c.doWriteIO()
}

func (c *Connection) writeBufLen() (bufLen int) {
    for _, buf := range c.writeBuffers {
        bufLen += len(buf)
    }
    return
}

func (c *Connection) doWriteIO() (bytesSent int64, err error) {
    buffers := c.writeBuffers
    bytesSent, err = buffers.WriteTo(c.rawConn)
    if err != nil {
        return bytesSent, err
    }

    for i, buf := range c.ioBuffers {
        c.ioBuffers[i] = nil
        c.writeBuffers[i] = nil
        if buf.EOF() {
            err = iobufferpool.ErrEOF
        }
        buf.Release()
    }
    c.ioBuffers = c.ioBuffers[:0]
    c.writeBuffers = c.writeBuffers[:0]
    return
}

// ServerConnection wrap tcp connection to backend server
type ServerConnection struct {
    Connection
    connectTimeout time.Duration
}

// NewServerConn construct tcp server connection
func NewServerConn(connectTimeout uint32, serverAddr net.Addr, listenerStopChan chan struct{}) *ServerConnection {
    conn := &ServerConnection{
        Connection: Connection{
            remoteAddr: serverAddr,

            writeBufferChan: make(chan *iobufferpool.StreamBuffer, writeBufSize),

            mu:               sync.Mutex{},
            connStopChan:     make(chan struct{}),
            listenerStopChan: listenerStopChan,
        },
        connectTimeout: time.Duration(connectTimeout) * time.Millisecond,
    }
    return conn
}

// Connect create backend server tcp connection
func (u *ServerConnection) Connect() bool {
    addr := u.remoteAddr
    if addr == nil {
        logger.Errorf("cannot connect because the server has been closed, server addr: %s", addr.String())
        return false
    }

    timeout := u.connectTimeout
    if timeout == 0 {
        timeout = 10 * time.Second
    }

    var err error
    u.rawConn, err = net.DialTimeout("tcp", addr.String(), timeout)
    if err != nil {
        if err == io.EOF {
            logger.Errorf("cannot connect because the server has been closed, server addr: %s", addr.String())
        } else if te, ok := err.(net.Error); ok && te.Timeout() {
            logger.Errorf("connect to server timeout, server addr: %s", addr.String())
        } else {
            logger.Errorf("connect to server failed, server addr: %s, err: %s", addr.String(), err.Error())
        }
        return false
    }

    u.localAddr = u.rawConn.LocalAddr()
    _ = u.rawConn.(*net.TCPConn).SetNoDelay(true)
    _ = u.rawConn.(*net.TCPConn).SetKeepAlive(true)
    u.Start()
    return true
}

以上是当初的实现，此段代码通过两个goroutine分别处理读写操作，并在读写过程中进行错误处理和流量控制。

    readBuffer      []byte
    writeBuffers    net.Buffers
    ioBuffers       []*iobufferpool.StreamBuffer
    writeBufferChan chan *iobufferpool.StreamBuffer

代码中还通过以上字段去控制读写缓冲区和写缓冲通道，虽然现在已经不清楚具体设计细节，但可以看出写的是相当不好，引入了太多复杂度，关键自己也没有很好的控制手段。

2 使用io.Copy

当初陈皓老师曾建议使用io.Copy来简化读写逻辑，因其已内置高效的流量拷贝和边界情况处理。
但当时我已按自行方案完成初版，既不舍推翻，也因工作紧张未能重构。
如今回看，采用io.Copy确实是更优选择，可借助Go标准库的成熟实现降低复杂度。

此外，UDP代理因其无连接特性带来较多实现复杂性，若从一开始就暂缓支持，或许更为合理。

2.1 简化后的TCP代理示例

package main

import (
    "io"
    "log"
    "net"
    "sync"
)

// TCPProxy 简单的TCP代理实现
type TCPProxy struct {
    ListenAddr string
    TargetAddr string
}

func (p *TCPProxy) Start() error {
    tcpAddr, err := net.ResolveTCPAddr("tcp", p.ListenAddr)
    if err != nil {
        return err
    }

    listener, err := net.ListenTCP("tcp", tcpAddr)
    if err != nil {
        return err
    }
    defer listener.Close()

    log.Printf("TCP代理启动，监听地址: %s，目标地址: %s", p.ListenAddr, p.TargetAddr)

    for {
        clientConn, err := listener.AcceptTCP()
        if err != nil {
            log.Printf("接受连接失败: %v", err)
            continue
        }

        go p.handleConnection(clientConn)
    }
}

func (p *TCPProxy) handleConnection(clientConn *net.TCPConn) {
    defer clientConn.Close()

    tcpAddr, _ := net.ResolveTCPAddr("tcp", p.TargetAddr)
    serverConn, _ := net.DialTCP("tcp", nil, tcpAddr)
    defer serverConn.Close()

    var wg sync.WaitGroup
    wg.Add(2)

    go func() {
        defer wg.Done()
        defer serverConn.Close()

        _, _ = io.Copy(serverConn, clientConn)
    }()

    go func() {
        defer wg.Done()
        defer clientConn.Close()

        _, _ = io.Copy(clientConn, serverConn)
    }()

    wg.Wait()
}

func main() {
    proxy := &TCPProxy{
        ListenAddr: ":8080",
        TargetAddr: "localhost:80", // 示例：代理到本地HTTP服务
    }

    if err := proxy.Start(); err != nil {
        log.Fatalf("代理启动失败: %v", err)
    }
}

2.2 添加缓冲区

简化的代码中，每次调用io.Copy均会使用默认缓冲区大小（32KB）进行数据传输。
此时，可能会想到使用io.CopyBuffer来指定更合适的缓冲区大小，以优化性能，代码示例如下：

func (p *TCPProxy) handleConnection(clientConn *net.TCPConn) {
    defer clientConn.Close()

    tcpAddr, _ := net.ResolveTCPAddr("tcp", p.TargetAddr)
    serverConn, _ := net.DialTCP("tcp", nil, tcpAddr)
    defer serverConn.Close()

    var wg sync.WaitGroup
    wg.Add(2)

    // 从客户端读取数据并写入服务器
    go func() {
        defer wg.Done()
        defer serverConn.CloseWrite()

        bufPtr := bufferPool.Get().(*[]byte)
        defer bufferPool.Put(bufPtr)

        _, _ = io.CopyBuffer(serverConn, clientConn, *bufPtr)
    }()

    // 从服务器读取数据并写入客户端
    go func() {
        defer wg.Done()
        defer clientConn.CloseWrite()

        bufPtr := bufferPool.Get().(*[]byte)
        defer bufferPool.Put(bufPtr)

        _, _ = io.CopyBuffer(clientConn, serverConn, *bufPtr)
    }()

    wg.Wait()
}

但是这样真的能提升性能吗？以下是io.CopyBuffer的实现代码：

// CopyBuffer is identical to Copy except that it stages through the
// provided buffer (if one is required) rather than allocating a
// temporary one. If buf is nil, one is allocated; otherwise if it has
// zero length, CopyBuffer panics.
//
// If either src implements [WriterTo] or dst implements [ReaderFrom],
// buf will not be used to perform the copy.
func CopyBuffer(dst Writer, src Reader, buf []byte) (written int64, err error) {
    if buf != nil && len(buf) == 0 {
        panic("empty buffer in CopyBuffer")
    }
    return copyBuffer(dst, src, buf)
}

// copyBuffer is the actual implementation of Copy and CopyBuffer.
// if buf is nil, one is allocated.
func copyBuffer(dst Writer, src Reader, buf []byte) (written int64, err error) {
    // If the reader has a WriteTo method, use it to do the copy.
    // Avoids an allocation and a copy.
    if wt, ok := src.(WriterTo); ok {
        return wt.WriteTo(dst)
    }
    // Similarly, if the writer has a ReadFrom method, use it to do the copy.
    if rf, ok := dst.(ReaderFrom); ok {
        return rf.ReadFrom(src)
    }
    if buf == nil {
        size := 32 * 1024
        if l, ok := src.(*LimitedReader); ok && int64(size) > l.N {
            if l.N < 1 {
                size = 1
            } else {
                size = int(l.N)
            }
        }
        buf = make([]byte, size)
    }
    for {
        nr, er := src.Read(buf)
        if nr > 0 {
            nw, ew := dst.Write(buf[0:nr])
            if nw < 0 || nr < nw {
                nw = 0
                if ew == nil {
                    ew = errInvalidWrite
                }
            }
            written += int64(nw)
            if ew != nil {
                err = ew
                break
            }
            if nr != nw {
                err = ErrShortWrite
                break
            }
        }
        if er != nil {
            if er != EOF {
                err = er
            }
            break
        }
    }
    return written, err
}

从代码可以看出，io.CopyBuffer在源或目标实现了WriterTo或ReaderFrom接口时，并不会使用提供的缓冲区。
而net.TCPConn类型正是实现了这些接口，具体代码详见：

// ReadFrom implements the [io.ReaderFrom] ReadFrom method.
func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
    if !c.ok() {
        return 0, syscall.EINVAL
    }
    n, err := c.readFrom(r)
    if err != nil && err != io.EOF {
        err = &OpError{Op: "readfrom", Net: c.fd.net, Source: c.fd.laddr, Addr: c.fd.raddr, Err: err}
    }
    return n, err
}

因此，在这种情况下，使用io.CopyBuffer并不会带来性能提升，反而是凭空添加了代码复杂度。

2.3 ReadFrom和WriteTo接口

net.TCPConn实现了io.ReaderFrom和io.WriterTo接口，这使得io.Copy能够直接利用这些高效的读写方法进行数据传输。
以下是ReadFrom和WriteTo的实现代码片段：

func (c *TCPConn) readFrom(r io.Reader) (int64, error) {
    if n, err, handled := spliceFrom(c.fd, r); handled {
        return n, err
    }
    if n, err, handled := sendFile(c.fd, r); handled {
        return n, err
    }
    return genericReadFrom(c, r)
}

func (c *TCPConn) writeTo(w io.Writer) (int64, error) {
    if n, err, handled := spliceTo(w, c.fd); handled {
        return n, err
    }
    return genericWriteTo(c, w)
}

查阅源码注释，可以发现spliceFrom、sendFile以及spliceTo等函数均利用了操作系统底层的高效数据传输机制，如Linux的splice和sendfile系统调用，从而实现零拷贝数据传输，极大提升了性能。

3 遗憾

在利用io.Copy简化数据流拷贝的同时，其对于长连接的空闲超时（Idle Timeout）管理能力确实存在不足。
*net.TCPConn提供的SetDeadline方法所设定的是连接级别的绝对超时时间，这对于HTTP短连接较为适用，但在需要维持连接长时间存活的四层代理场景中则缺乏灵活性。
一个可行的优化思路是，封装支持超时控制的Reader/Writer接口，在每次读写操作前动态更新超时时间，从而实现真正的空闲超时机制。
然而，这种方案需要自行实现完整的读写循环，无法直接复用io.Copy等系统级API，这在一定程度上又回到了最初的复杂度问题，需要在简洁性与控制粒度之间进行权衡。

还是不折腾了，寄希望于Go未来在net包中提供更灵活的超时控制机制吧，正如Golang的.ioCopy引入了ReadFrom和WriteTo接口一样。