64行代码实现零拷贝go的TCP拆包粘包

前言

这段时间想用go写一个简单IM系统，就思考了一下go语言TCP的拆包粘包。TCP的拆包粘包有一般有三种解决方案。

使用定长字节

实际使用中，少于固定字长的，要用字符去填充，空间使用率不够高。

使用分隔符

一般用文本传输的，使用分隔符，IM系统一般对性能要求高，不推荐使用文本传输。

用消息的头字节标识消息内容的长度

可以使用二进制传输，效率高，推荐。下面看看怎么实现。

尝试使用系统库自带的bytes.Buffer实现

代码实现：

package tcp

import (
    "fmt"
    "net"
    "log"
    "bytes"
    "encoding/binary"
)

const (
    BYTES_SIZE uint16 = 1024
    HEAD_SIZE  int    = 2
)

func StartServer(address string) {
    listener, err := net.Listen("tcp", address)
    if err != nil {
        log.Println("Error listening", err.Error())
        return
    }
    for {
        conn, err := listener.Accept()
        fmt.Println(conn.RemoteAddr())
        if err != nil {
            fmt.Println("Error accepting", err.Error())
            return // 终止程序
        }
        go doConn(conn)
    }
}

func doConn(conn net.Conn) {
    var (
        buffer           = bytes.NewBuffer(make([]byte, 0, BYTES_SIZE))
        bytes            = make([]byte, BYTES_SIZE);
        isHead      bool = true
        contentSize int
        head        = make([]byte, HEAD_SIZE)
        content     = make([]byte, BYTES_SIZE)
    )
    for {
        readLen, err := conn.Read(bytes);
        if err != nil {
            log.Println("Error reading", err.Error())
            return
        }
        _, err = buffer.Write(bytes[0:readLen])
        if err != nil {
            log.Println("Error writing to buffer", err.Error())
            return
        }

        for {
            if isHead {
                if buffer.Len() >= HEAD_SIZE {
                    _, err := buffer.Read(head)
                    if err != nil {
                        fmt.Println("Error reading", err.Error())
                        return
                    }
                    contentSize = int(binary.BigEndian.Uint16(head))
                    isHead = false
                } else {
                    break
                }
            }
            if !isHead {
                if buffer.Len() >= contentSize {
                    _, err := buffer.Read(content[:contentSize])
                    if err != nil {
                        fmt.Println("Error reading", err.Error())
                        return
                    }
                    fmt.Println(string(content[:contentSize]))
                    isHead = true
                } else {
                    break
                }
            }
        }
    }
}

测试用例：

package tcp

import (
    "testing"
    "net"
    "fmt"
    "encoding/binary"
)

func TestStartServer(t *testing.T) {
    StartServer("localhost:50002")
}

func TestClient(t *testing.T) {
    conn, err := net.Dial("tcp", "localhost:50002")
    if err != nil {
        fmt.Println("Error dialing", err.Error())
        return // 终止程序
    }
    var headSize int
    var headBytes = make([]byte, 2)
    s := "hello world"
    content := []byte(s)
    headSize = len(content)
    binary.BigEndian.PutUint16(headBytes, uint16(headSize))
    conn.Write(headBytes)
    conn.Write(content)

    s = "hello go"
    content = []byte(s)
    headSize = len(content)
    binary.BigEndian.PutUint16(headBytes, uint16(headSize))
    conn.Write(headBytes)
    conn.Write(content)

    s = "hello tcp"
    content = []byte(s)
    headSize = len(content)
    binary.BigEndian.PutUint16(headBytes, uint16(headSize))
    conn.Write(headBytes)
    conn.Write(content)
}

执行结果

127.0.0.1:51062
hello world
hello go
hello tcp

用go系统库的buffer，是不是感觉代码特别别扭，两大缺点

1.要写大量的逻辑代码，来弥补buffer对这个场景的不适用。

2.性能不高，有三次次内存拷贝，coon->[]byte->Buffer->[]byte。

自己实现

既然轮子不合适，就自己造轮子，首先实现一个自己的Buffer,很简单，只有六十几行代码，所有过程只有一次byte数组的拷贝，conn->buffer,剩下的全部操作都在原buffer的字节数组里面操作

package tcp

import (
    "errors"
    "io"
)

type buffer struct {
    reader io.Reader
    buf    []byte
    start  int
    end    int
}

func newBuffer(reader io.Reader, len int) buffer {
    buf := make([]byte, len)
    return buffer{reader, buf, 0, 0}
}

func (b *buffer) Len() int {
    return b.end - b.start
}

//将有用的字节前移
func (b *buffer) grow() {
    if b.start == 0 {
        return
    }
    copy(b.buf, b.buf[b.start:b.end])
    b.end -= b.start
    b.start = 0;
}

//从reader里面读取数据，如果reader阻塞，会发生阻塞
func (b *buffer) readFromReader() (int, error) {
    b.grow()
    n, err := b.reader.Read(b.buf[b.end:])
    if (err != nil) {
        return n, err
    }
    b.end += n
    return n, nil
}

//返回n个字节，而不产生移位
func (b *buffer) seek(n int) ([]byte, error) {
    if b.end-b.start >= n {
        buf := b.buf[b.start:b.start+n]
        return buf, nil
    }
    return nil, errors.New("not enough")
}

//舍弃offset个字段，读取n个字段
func (b *buffer) read(offset, n int) ([]byte) {
    b.start += offset
    buf := b.buf[b.start:b.start+n]
    b.start += n
    return buf
}

再看看怎样使用它,将上面的doConn函数改成这样就行了。

func doConn(conn net.Conn) {
    var (
        buffer      = newBuffer(conn, 16)
        headBuf     []byte
        contentSize int
        contentBuf  []byte
    )
    for {
        _, err := buffer.readFromReader()
        if err != nil {
            fmt.Println(err)
            return
        }
        for {
            headBuf, err = buffer.seek(HEAD_SIZE);
            if err != nil {
                break
            }
            contentSize = int(binary.BigEndian.Uint16(headBuf))
            if (buffer.Len() >= contentSize-HEAD_SIZE) {
                contentBuf = buffer.read(HEAD_SIZE, contentSize)
                fmt.Println(string(contentBuf))
                continue
            }
            break
        }
    }
}

跑下测试用例，看下结果

127.0.0.1:51062
hello world
hello go
hello tcp

源码地址：https://github.com/alberliu/goim

你有更好的方式，可以邮箱我，alber_liu@qq.com,让我学习一下