用Go写一个简单的TCP server or client 模型

对Conn封装的基本思路

go内置了net包已经很好的封装了socket通讯。然而在实际使用中，由于net/Conn的Read/Write方法是堵塞的原因，必须将其放入单独的goroutine中进行处理。

我们先简单的整理下思路，对于连接(Conn)的处理，我们可以开启2条goroutine进行处理，一条用于堵塞的Read的处理，另一条进行Write的处理。

这里必须指出，其实Write本身就是线程安全的，也就是我们在别的任何地方都可以进行Write，但是Write是堵塞的，所以考虑到这点，我们还是将其放入一个单独的goroutine中进行处理。

这样设计的原因在于Conn是支持同时Read/Write的，这样我们的基本的Conn的模型就成型了。对于服务端或者客户端而言，我们只需要封装对应过来的Conn即可，Conn的读写goroutine进行处理，并将获得的事件抛向外部。

那么我们就按这个思路来实现一个简单的Connection封装，该封装支持线程安全的写，并且支持解包操作。

package tcpnetwork

import (
"errors"
"log"
"net"
"time"
)

const (
kConnStatus_None = iota
kConnStatus_Connected
kConnStatus_Disconnected
)

const (
kConnEvent_None = iota
kConnEvent_Connected
kConnEvent_Disconnected
kConnEvent_Data
kConnEvent_Close
)

const (
kConnConf_DefaultSendTimeoutSec = 5
kConnConf_MaxReadBufferLength   = 0xffff // 0xffff
)

type Connection struct {
conn                net.Conn
status              int
connId              int
sendMsgQueue        chan []byte
sendTimeoutSec      int
eventQueue          IEventQueue
streamProtocol      IStreamProtocol
maxReadBufferLength int
userdata            interface{}
from                int
readTimeoutSec      int
}

func newConnection(c net.Conn, sendBufferSize int, eq IEventQueue) *Connection {
return &Connection{
conn:                c,
status:              kConnStatus_None,
connId:              0,
sendMsgQueue:        make(chan []byte, sendBufferSize),
sendTimeoutSec:      kConnConf_DefaultSendTimeoutSec,
maxReadBufferLength: kConnConf_MaxReadBufferLength,
eventQueue:          eq,
}
}

type ConnEvent struct {
EventType int
Conn      *Connection
Data      []byte
}

func newConnEvent(et int, c *Connection, d []byte) *ConnEvent {
return &ConnEvent{
EventType: et,
Conn:      c,
Data:      d,
}
}

//  directly close, packages in queue will not be sent
func (this *Connection) close() {
if kConnStatus_Connected != this.status {
return
}

this.conn.Close()
this.status = kConnStatus_Disconnected
}

func (this *Connection) Close() {
if this.status != kConnStatus_Connected {
return
}

select {
case this.sendMsgQueue <- nil:
{
//  nothing
}
case <-time.After(time.Duration(this.sendTimeoutSec)):
{
//  timeout, close the connection
this.close()
log.Printf("Con[%d] send message timeout, close it", this.connId)
}
}
}

func (this *Connection) pushEvent(et int, d []byte) {
if nil == this.eventQueue {
log.Println("Nil event queue")
return
}
this.eventQueue.Push(newConnEvent(et, this, d))
}

func (this *Connection) GetStatus() int {
return this.status
}

func (this *Connection) setStatus(stat int) {
this.status = stat
}

func (this *Connection) GetConnId() int {
return this.connId
}

func (this *Connection) SetConnId(id int) {
this.connId = id
}

func (this *Connection) GetUserdata() interface{} {
return this.userdata
}

func (this *Connection) SetUserdata(ud interface{}) {
this.userdata = ud
}

func (this *Connection) SetReadTimeoutSec(sec int) {
this.readTimeoutSec = sec
}

func (this *Connection) GetReadTimeoutSec() int {
return this.readTimeoutSec
}

func (this *Connection) setStreamProtocol(sp IStreamProtocol) {
this.streamProtocol = sp
}

func (this *Connection) sendRaw(msg []byte) {
if this.status != kConnStatus_Connected {
return
}

select {
case this.sendMsgQueue <- msg:
{
//  nothing
}
case <-time.After(time.Duration(this.sendTimeoutSec)):
{
//  timeout, close the connection
this.close()
log.Printf("Con[%d] send message timeout, close it", this.connId)
}
}
}

func (this *Connection) Send(msg []byte, cpy bool) {
if this.status != kConnStatus_Connected {
return
}

buf := msg
if cpy {
msgCopy := make([]byte, len(msg))
copy(msgCopy, msg)
buf = msgCopy
}

select {
case this.sendMsgQueue <- buf:
{
//  nothing
}
case <-time.After(time.Duration(this.sendTimeoutSec)):
{
//  timeout, close the connection
this.close()
log.Printf("Con[%d] send message timeout, close it", this.connId)
}
}
}

//  run a routine to process the connection
func (this *Connection) run() {
go this.routineMain()
}

func (this *Connection) routineMain() {
defer func() {
//  routine end
log.Printf("Routine of connection[%d] quit", this.connId)
e := recover()
if e != nil {
log.Println("Panic:", e)
}

//  close the connection
this.close()

//  free channel
close(this.sendMsgQueue)
this.sendMsgQueue = nil

//  post event
this.pushEvent(kConnEvent_Disconnected, nil)
}()

if nil == this.streamProtocol {
log.Println("Nil stream protocol")
return
}
this.streamProtocol.Init()

//  connected
this.pushEvent(kConnEvent_Connected, nil)
this.status = kConnStatus_Connected

go this.routineSend()
this.routineRead()
}

func (this *Connection) routineSend() error {
defer func() {
log.Println("Connection", this.connId, " send loop return")
}()

for {
select {
case evt, ok := <-this.sendMsgQueue:
{
if !ok {
//  channel closed, quit
return nil
}

if nil == evt {
log.Println("User disconnect")
this.close()
return nil
}

var err error

headerBytes := this.streamProtocol.SerializeHeader(evt)
if nil != headerBytes {
//  write header first
_, err = this.conn.Write(headerBytes)
if err != nil {
log.Println("Conn write error:", err)
return err
}
}

_, err = this.conn.Write(evt)
if err != nil {
log.Println("Conn write error:", err)
return err
}
}
}
}

return nil
}

func (this *Connection) routineRead() error {
//  default buffer
buf := make([]byte, this.maxReadBufferLength)

for {
msg, err := this.unpack(buf)
if err != nil {
log.Println("Conn read error:", err)
return err
}

this.pushEvent(kConnEvent_Data, msg)
}

return nil
}

func (this *Connection) unpack(buf []byte) ([]byte, error) {
//  read head
if 0 != this.readTimeoutSec {
this.conn.SetReadDeadline(time.Now().Add(time.Duration(this.readTimeoutSec) * time.Second))
}
headBuf := buf[:this.streamProtocol.GetHeaderLength()]
_, err := this.conn.Read(headBuf)
if err != nil {
return nil, err
}

//  check length
packetLength := this.streamProtocol.UnserializeHeader(headBuf)
if packetLength > this.maxReadBufferLength ||
0 == packetLength {
return nil, errors.New("The stream data is too long")
}

//  read body
if 0 != this.readTimeoutSec {
this.conn.SetReadDeadline(time.Now().Add(time.Duration(this.readTimeoutSec) * time.Second))
}
bodyLength := packetLength - this.streamProtocol.GetHeaderLength()
_, err = this.conn.Read(buf[:bodyLength])
if err != nil {
return nil, err
}

//  ok
msg := make([]byte, bodyLength)
copy(msg, buf[:bodyLength])
if 0 != this.readTimeoutSec {
this.conn.SetReadDeadline(time.Time{})
}

return msg, nil
}

这就是简单的Conn封装，在抛出事件这部，定义了一个interface来接收事件。

package tcpnetwork

type IEventQueue interface {
Push(*ConnEvent)
Pop() *ConnEvent
}

type IStreamProtocol interface {
//  Init
Init()
//  get the header length of the stream
GetHeaderLength() int
//  read the header length of the stream
UnserializeHeader([]byte) int
//  format header
SerializeHeader([]byte) []byte
}

type IEventHandler interface {
OnConnected(evt *ConnEvent)
OnDisconnected(evt *ConnEvent)
OnRecv(evt *ConnEvent)
}

我们只要在实现对应的方法，就可以接收事件和读取事件了。

Server/Client 端的实现

Server

我们已经封装好了Conn，那么接下来的工作将会简单很多。我们来封装一个TCPNetwork的结构。

对于Server端来说，它基本的步骤就是

监听端口

开启accept线程

得到Conn，并生成Connection，开启Read/Write线程

第一步很简单，net包已封装

ls, err := net.Listen("tcp", addr)
if nil != err {
return err
}

//  accept
this.listener = ls
go this.acceptRoutine()
return nil

我们在listen成功后，开启了一个goroutine来不断的进行死循环来等待连接的接入。

func (this *TCPNetwork) acceptRoutine() {
for {
conn, err := this.listener.Accept()
if err != nil {
log.Println("accept routine quit.error:", err)
return
}

//  process conn event
connection := this.createConn(conn)
connection.SetReadTimeoutSec(this.readTimeoutSec)
connection.from = 0
connection.run()
}
}

得到了Connection后，我们只需要让它的处理routine跑起来即可。然而我们需要对应的Connection抛过来的事件，于是我们在TCPNetwork中实现IEventQueue的2个方法，这样我们的TCPNetwork就可以接收对应的事件的抛入，也可以读取，大家读到这里也就知道了，最适合实现这个的就是golang的神器：channel。

我们来为TCPNetwork定义一个chan *ConnEvent，并实现接口的方法。

func (this *TCPNetwork) Push(evt *ConnEvent) {
if nil == this.eventQueue {
return
}
this.eventQueue <- evt
}

func (this *TCPNetwork) Pop() *ConnEvent {
evt, ok := <-this.eventQueue
if !ok {
//  event queue already closed
this.eventQueue = nil
return nil
}

return evt
}

其实基本的逻辑已经完成了，对于我们来说，只需要关心eventQueue里的内容就行了，这就属于上层逻辑处理了。这样一个简单的TCPNetwork就封装好了。

Client

Client端基本没什么好说的，net包Connect成功后会获得一个Conn，然后对于这个Conn的处理其实和Server端一样了。

func (this *TCPNetwork) Connect(addr string) error {
conn, err := net.Dial("tcp", addr)
if nil != err {
return err
}

connection := this.createConn(conn)
connection.from = 1
connection.run()

return nil
}

使用方法

对于外部来说，使用很简单，这里贴上一个简单的echo server example.

package main

import (
"log"

"github.com/sryanyuan/tcpnetwork"
)

type TSShutdown struct {
network *tcpnetwork.TCPNetwork
}

func NewTSShutdown() *TSShutdown {
t := &TSShutdown{}
t.network = tcpnetwork.NewTCPNetwork(1024, tcpnetwork.NewStreamProtocol4())
return t
}

func (this *TSShutdown) OnConnected(evt *tcpnetwork.ConnEvent) {
log.Println("connected ", evt.Conn.GetConnId())
}

func (this *TSShutdown) OnDisconnected(evt *tcpnetwork.ConnEvent) {
log.Println("disconnected ", evt.Conn.GetConnId())
}

func (this *TSShutdown) OnRecv(evt *tcpnetwork.ConnEvent) {
log.Println("recv ", evt.Conn.GetConnId(), evt.Data)

evt.Conn.Send(evt.Data, false)
}

func (this *TSShutdown) Run() {
err := this.network.Listen("127.0.0.1:2222")

if err != nil {
log.Println(err)
return
}

this.network.ServeWithHandler(this)
log.Println("done")
}

package main

import (
"fmt"
"log"
)

func main() {
defer func() {
e := recover()
if e != nil {
log.Println(e)
}
var inp int
fmt.Scanln(&inp)
}()
tsshutdown := NewTSShutdown()
tsshutdown.Run()
}

总结

我们其实已经实现了一个简单的tcp封装了，支持server/client连接，至于其中心跳的细节、封包解包的细节等等，这里就不多介绍了，可以通过阅读源码来理解。

该封装可以见我的Github TCPNetwork