打造先进的内存KV数据库-5 TCP侦听

TCP侦听

作为支持集群的数据库，必定要与多个客户端交互信息，不可能让数据库与所有客户共享地址空间（虽然这样性能好），所以需要使用TCP协议进行交互数据，（UDP协议不可靠。。。弃用），C语言的TCP库其实还好，但是对于高并发和并行的处理不如Go，而且并发锁机制比较难写，所以使用Go写了服务器和客户端调用C的库，目前版本没有什么身份验证，之后会加上。

代码实现

//server.go
package main
// #cgo LDFLAGS: -L ./lib -lmonkeyS
// #include "./lib/core.h"
// #include <stdlib.h>
import "C"
import (
    "unsafe"
    _"fmt"
    "net"
    "strings"
)

func main() {
    str := []byte("monkey")
    str = append(str,0)
    C.CreateDB((*C.char)(unsafe.Pointer(&str[0])))  //创建基础数据库
    servicePort := ":1517"
    tcpAddr,err := net.ResolveTCPAddr("tcp4",servicePort)
    if err != nil {
        panic(err)
    }
    l,err := net.ListenTCP("tcp",tcpAddr)   //侦听TCP
    if err != nil {
        panic(err)
    }
    for{
        conn,err := l.Accept()
        if err != nil {
            panic(err)
        }
        go Handler(conn)
    }
}

func Handler(conn net.Conn) {

    str := []byte("monkey")                         //环境变量-当前数据库
    db := C.SwitchDB((*C.char)(unsafe.Pointer(&str[0])))
    for {               
        buff := []byte{}
        buf := make([]byte,1024)
        length,err := conn.Read(buf)
        total := uint32(0); //前4个字节保存消息长度
        for i := 0;i < 4;i++ {
            total <<= 8;
            total += uint32(buf[i]);
        }
        //fmt.Println("Message length:",total)
        buff = append(buff,buf[4:]...)
        total -= uint32(length)
        for total > 0 {
            length,err = conn.Read(buf)
            total -= uint32(length)
            buff = append(buff,buf...)
        }
        if err != nil {
            conn.Close()
            break
        }
        TranslateMessage(conn,&db,buff)                     //解析消息
    }

}

func TranslateMessage(conn net.Conn,db **C.Database,message []byte) {
    command := string(message)
    params := strings.Split(command," ")
    //fmt.Println(params)
    response := []byte{}
    if params[0] == "set" {
        r := C.Set(&(*db).tIndex,(*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))),(unsafe.Pointer(&([]byte(params[2]))[0])))
        for i := 0;;i++ {
            response = append(response,byte(r.msg[i]))
            if response[i] == 0 { break; }
        }

    }else if params[0] == "get" {
        r := C.Get(&(*db).tIndex,(*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))))
        // for i := 0;;i++ {
        //  response = append(response,byte(r.msg[i]))
        //  if response[i] == 0 { break; }
        // }
        if(int(r.code) == 0) {
            for i := 0;;i++ {
                response = append(response,byte(*(*C.char)(unsafe.Pointer((uintptr(r.pData)+uintptr(i))))))
                if response[i] == 0 { break; }
            }
        }else {
            // for i := 0;;i++ {
            // response = append(response,byte(r.msg[i]))
            // if response[i] == 0 { break; }
            // }
        }

    }else if params[0] == "delete" || params[0] == "remove" {
        r := C.Delete(&(*db).tIndex,(*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))))
        for i := 0;;i++ {
            response = append(response,byte(r.msg[i]))
            if response[i] == 0 { break; }
        }

    }else if params[0] == "createdb" {
        d := C.CreateDB((*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))))
        if d != nil {
            *db = d
            response = []byte("Already exist,switched\n")
        }else {
            response = []byte("Created\n")
        }
    }else if params[0] == "switchdb" {
        d := C.SwitchDB((*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))))
        if d != nil {
            *db = d
            response = []byte("ok\n")
        }else {
            response = []byte("fail\n")
        }
    }else if params[0] == "dropdb" {
        *db = C.DropDB((*C.char)(unsafe.Pointer((&([]byte(params[1]))[0]))))
    }else if strings.EqualFold("listdb",params[0]) {
        r := C.ListDB()
        for i := 0;i < 1024;i++ {
            b := byte(*(*C.char)(unsafe.Pointer(uintptr(unsafe.Pointer(r))+uintptr(i))))
            response = append(response,b)
            if(b == 0){ break; }
        }
        C.free(unsafe.Pointer(r))
    }else {
        //fmt.Println("unkown command:",params[0])
    }
    total := len(response) + 4
    header := make([]byte,4)
    i := 0
    for total > 0 {
        header[3-i] = byte(total % 256)
        total /= 256
        i++
    }
    response = append(header,response...)
    conn.Write(response)
}

//Client.go
package main
import "net"
import "fmt"
func main() {
    tcpAddr, err := net.ResolveTCPAddr("tcp4", "127.0.0.1:1517")  
    if err != nil {
        panic(err)
    }
    conn, err := net.DialTCP("tcp", nil, tcpAddr)  
    if err != nil {
        panic(err)
    }

    for {
        buf1 := ""
        buf2 := ""
        buf3 := ""
        buf := ""
        fmt.Print("monkey>")
        fmt.Scanf("%s",&buf1)
        if buf1 == "set" {
            fmt.Scanf("%s",&buf2)
            fmt.Scanf("%s",&buf3)
            buf = buf1 + " " + buf2 + " " + buf3
        }else if buf1 == "get"{
            fmt.Scanf("%s",&buf2)
            buf = buf1 + " " + buf2
        }else if buf1 == "remove" || buf1 == "delete" {
            fmt.Scanf("%s",&buf2)
            buf = buf1 + " " + buf2
        }else if buf1 == "createdb"{
            fmt.Scanf("%s",&buf2)
            buf = buf1 + " " + buf2
        }else if buf1 == "switchdb"{
            fmt.Scanf("%s",&buf2)
            buf = buf1 + " " + buf2
        }else if buf1 == "dropdb"{
            fmt.Scanf("%s",&buf2)
            buf = buf1 + " " + buf2
        }else if buf1 == "listdb"{
            buf = buf1 + " "
        }else if buf1 == "exit"{
            fmt.Println("Bye!")
            break;
        }
        total := uint32(0)
        total = uint32(len(buf) + 4)
        header := make([]byte,4)
        i := 0
        for total > 0 {
            header[3-i] = byte(total % 256)
            total /= 256
            i++
        } 
        conn.Write(append(header,([]byte(buf))...))

        buff := []byte{}
        buff2 := make([]byte,1024)
        length,_ := conn.Read(buff2)
        total = uint32(0);  //前4个字节保存消息长度
        for i := 0;i < 4;i++ {
            total <<= 8;
            total += uint32(buff2[i]);
        }
        buff = append(buff,buff2[4:]...)
        total -= uint32(length)
        for total > 0 {
            length,_ = conn.Read(buff2)
            total -= uint32(length)
            buff = append(buff,buff2...)
        }
        for i := 0;i < 1024;i++ {
            if buff[i] == 0 { break; }
            fmt.Printf("%c",buff[i])
        }
        fmt.Print("\n")
    }
}

修正：上述代码存在严重问题：
发送1K以上数据会无法正确接收
改进代码如下：

//tcp.go
package tcp
import "net"
import "fmt"

func ok(bytes []byte) bool {
    return bytes[0] == 111 && bytes[1] == 107 && bytes[2] == 0;
}

func bytes4uint(bytes []byte) uint32 {
    total := uint32(0); 
    for i := 0;i < 4;i++ {
        total <<= 8;
        total += uint32(bytes[i]);
    }
    return total
}

func uint32bytes(n uint32) []byte {
    header := make([]byte,4)
    i := 0
    for n > 0 {
        header[3-i] = byte(n % 256)
        n /= 256
        i++
    }
    return header
}


type TCPSession struct {
    Conn *net.TCPConn
    ToSend chan interface{} //要发送的数据
    Received chan interface{}   //接受到的数据
    Closed bool //是否已经关闭
}

func (s *TCPSession) Init() {
    s.ToSend = make(chan interface{})
    s.Received = make(chan interface{})
    go s.Send()
    go s.Recv()
}

func (s *TCPSession) Send() {
    for {
        if s.Closed {
            return
        }
        buf0 := <- s.ToSend //取出要发送的数据
        buf := buf0.([]byte)

        _,err := s.Conn.Write(buf)  //发送掉   
        //fmt.Println("send,",buf)
        if err != nil {
            s.Closed = true
            return
        }
    }

}

func (s *TCPSession) Recv() {
    for {
        if s.Closed {
            return
        }
        buf := make([]byte,1024)
        _,err := s.Conn.Read(buf)
        if err != nil {
            s.Closed = true
            return
        }
        s.Received <- buf
        //fmt.Println("read,",buf)
        }

}

func (s *TCPSession) SendMessage(bytes []byte) {
    total := len(bytes) / 1024
    if len(bytes) % 1024 != 0 {
        total++
    }
    header := uint32bytes(uint32(total))    //计算条数
    s.ToSend <- header
    //fmt.Println(header)
    for i := 0;i < total-1;i++ {
        buf := bytes[0:1024]    //发送这一段
        bytes = bytes[1024:]
        s.ToSend <- buf
        continue
    }
    //发送最后一段
    if total == 0 {
        return
    }
    buf := bytes[0:]    //发送这一段
    s.ToSend <- buf
}

func (s *TCPSession) ReadMessage() []byte {
    buf0 := <- s.Received
    buf := buf0.([]byte)
    //fmt.Println(buf)
    total := bytes4uint(buf)
    var buff []byte
    if buf[4] != 0 {    //两份报表被合并
        buff = buf[4:]
        total--
    } else {
        buff = []byte{}     
    }

    for i := uint32(0);i < total;i++ {
        buf0 := <- s.Received
        buf := buf0.([]byte)
        buff = append(buff,buf...)
    }
    return buff
}