tcpsock for Golang
Posted ecofast
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了tcpsock for Golang相关的知识，希望对你有一定的参考价值。
前记：本文所述的 tcpsock 库托管在 Github。
Golang 中的 net 标准库已对 TCP 网络编程作了简洁（却很不简单）的封装，基本上，可直接通过引用其提供的相关接口开发简易的网络应用程序。但我想封装一套简单并提供简洁调用接口的 TCP 网络库（即 tcpsock），以达到所谓组件化的目的，如此用户只需调用几个非常简单的接口（包括设置回调函数等）就能较直观地处理网络收发等逻辑。
一开始设计 tcpsock 时，只打算提供 TcpServer 及 TcpConn 对象而没有考虑加入 TcpClient（毕竟 Golang 多用于后台开发），但譬如网络游戏中的网关（GameGate），其既是服务器程序（服务于数计的网络客户端），又要充当客户端连接到游戏服务器（GameServer），此时即凸显 TcpClient 存在的必要了。所以重构了部分结构与代码，加入了 TcpClient 对象，而这也带来了额外的好处，具体可参考 chatroom 示例，其服务器与客户端在诸如协议解析等方面完全复用代码，大大减少了开发维护成本。
囿于个人当前的能力及精力， tcpsock 不可避免会存在一些甚至许多问题或 Bug（虽然它确实很简单），希望能在实际的项目运用及与技术同仁的探讨中一点点改进它，直至完善。
tcpsock 库共 4 个单元（tcpsock、tcpserver、tcpconn 及 tcpclient），并提供 3 个导出类型（TcpServer、TcpClient 及 TcpConn，分别实现于与其名称对应的单元），各单元代码如下：
// Copyright (C) 2017 ecofast(胡光耀). All rights reserved.
// Use of this source code is governed by a BSD-style license.
// Package tcpsock provides easy to use interfaces for TCP I/O.
// Thanks to darksword(gansidui) and AlexStocks for their valuable projects
// which are gotcp(https://github.com/gansidui/gotcp)
// and getty(https://github.com/AlexStocks/getty).
package tcpsock
import (
"sync"
)
const (
RecvBufLenMax = 4 * 1024
SendBufLenMax = 4 * 1024
SendBufCapMax = 10
RecvBufCapMax = 10
)
type tcpSock struct {
sendBufCap uint32
recvBufCap uint32
proto Protocol
exitChan chan struct{}
waitGroup *sync.WaitGroup
onConnConnect OnTcpConnCallback
onConnClose OnTcpConnCallback
}
type Protocol interface {
Parse(b []byte, recvChan chan<- Packet)
Process(conn *TcpConn, p Packet)
}
type Packet interface {
Marshal() []byte
}
// Copyright (C) 2017 ecofast(胡光耀). All rights reserved.
// Use of this source code is governed by a BSD-style license.
package tcpsock
import (
"net"
"sync"
"sync/atomic"
"time"
. "github.com/ecofast/sysutils"
)
type TcpServer struct {
listener *net.TCPListener
acceptTimeout int
*tcpSock
autoIncID uint32
numOfConn uint32
}
func NewTcpServer(listenPort, acceptTimeout int, protocol Protocol) *TcpServer {
tcpAddr, err := net.ResolveTCPAddr("tcp", ":"+IntToStr(int(listenPort)))
CheckError(err)
listener, err := net.ListenTCP("tcp", tcpAddr)
CheckError(err)
return &TcpServer{
listener: listener,
acceptTimeout: acceptTimeout,
tcpSock: &tcpSock{
sendBufCap: SendBufCapMax,
recvBufCap: RecvBufCapMax,
proto: protocol,
exitChan: make(chan struct{}),
waitGroup: &sync.WaitGroup{},
},
}
}
func (self *TcpServer) Serve() {
self.waitGroup.Add(1)
defer func() {
self.listener.Close()
self.waitGroup.Done()
}()
for {
select {
case <-self.exitChan:
return
default:
}
self.listener.SetDeadline(time.Now().Add(time.Duration(self.acceptTimeout) * time.Second))
conn, err := self.listener.AcceptTCP()
if err != nil {
continue
}
atomic.AddUint32(&self.numOfConn, 1)
self.waitGroup.Add(1)
go func() {
c := newTcpConn(atomic.AddUint32(&self.autoIncID, 1), self.tcpSock, conn, self.sendBufCap, self.recvBufCap, self.connClose)
if self.onConnConnect != nil {
self.onConnConnect(c)
}
c.run()
self.waitGroup.Done()
}()
}
}
func (self *TcpServer) Close() {
close(self.exitChan)
self.waitGroup.Wait()
}
func (self *TcpServer) NumOfConn() uint32 {
return atomic.LoadUint32(&self.numOfConn)
}
func (self *TcpServer) connClose(conn *TcpConn) {
atomic.AddUint32(&self.numOfConn, ^uint32(0))
if self.onConnClose != nil {
self.onConnClose(conn)
}
}
func (self *TcpServer) OnConnConnect(fn OnTcpConnCallback) {
self.onConnConnect = fn
}
func (self *TcpServer) OnConnClose(fn OnTcpConnCallback) {
self.onConnClose = fn
}
// Copyright (C) 2017 ecofast(胡光耀). All rights reserved.
// Use of this source code is governed by a BSD-style license.
package tcpsock
import (
"net"
"sync"
"sync/atomic"
)
type OnTcpConnCallback func(c *TcpConn)
type TcpConn struct {
id uint32
owner *tcpSock
conn *net.TCPConn
sendChan chan Packet
recvChan chan Packet
closeChan chan struct{}
closeOnce sync.Once
closedFlag int32
onClose OnTcpConnCallback
}
func newTcpConn(id uint32, owner *tcpSock, conn *net.TCPConn, sendCap, recvCap uint32, onClose OnTcpConnCallback) *TcpConn {
return &TcpConn{
id: id,
owner: owner,
conn: conn,
sendChan: make(chan Packet, sendCap),
recvChan: make(chan Packet, recvCap),
closeChan: make(chan struct{}),
onClose: onClose,
}
}
func (self *TcpConn) ID() uint32 {
return self.id
}
func (self *TcpConn) run() {
startGoroutine(self.reader, self.owner.waitGroup)
startGoroutine(self.writer, self.owner.waitGroup)
startGoroutine(self.handler, self.owner.waitGroup)
}
func (self *TcpConn) Close() {
self.closeOnce.Do(func() {
atomic.StoreInt32(&self.closedFlag, 1)
close(self.sendChan)
close(self.recvChan)
close(self.closeChan)
self.conn.Close()
if self.onClose != nil {
self.onClose(self)
}
})
}
func (self *TcpConn) Closed() bool {
return atomic.LoadInt32(&self.closedFlag) == 1
}
func (self *TcpConn) RawConn() *net.TCPConn {
return self.conn
}
func startGoroutine(fn func(), wg *sync.WaitGroup) {
wg.Add(1)
go func() {
fn()
wg.Done()
}()
}
func (self *TcpConn) reader() {
defer func() {
recover()
self.Close()
}()
buf := make([]byte, RecvBufLenMax)
for {
select {
case <-self.owner.exitChan:
return
case <-self.closeChan:
return
default:
}
count, err := self.conn.Read(buf)
if err != nil {
return
}
self.owner.proto.Parse(buf[:count], self.recvChan)
}
}
func (self *TcpConn) writer() {
defer func() {
recover()
self.Close()
}()
for {
if self.Closed() {
return
}
select {
case <-self.owner.exitChan:
return
case <-self.closeChan:
return
case p := <-self.sendChan:
if _, err := self.conn.Write(p.Marshal()); err != nil {
return
}
}
}
}
func (self *TcpConn) handler() {
defer func() {
recover()
self.Close()
}()
for {
if self.Closed() {
return
}
select {
case <-self.owner.exitChan:
return
case <-self.closeChan:
return
case packet := <-self.recvChan:
self.owner.proto.Process(self, packet)
}
}
}
func (self *TcpConn) Write(p Packet) {
if self.Closed() {
return
}
defer func() {
recover()
}()
self.sendChan <- p
}
// Copyright (C) 2017 ecofast(胡光耀). All rights reserved.
// Use of this source code is governed by a BSD-style license.
package tcpsock
import (
"net"
"sync"
. "github.com/ecofast/sysutils"
)
type TcpClient struct {
svrAddr *net.TCPAddr
*tcpSock
}
func NewTcpClient(svrAddr string, proto Protocol) *TcpClient {
tcpAddr, err := net.ResolveTCPAddr("tcp", svrAddr)
CheckError(err)
return &TcpClient{
svrAddr: tcpAddr,
tcpSock: &tcpSock{
sendBufCap: SendBufCapMax,
recvBufCap: RecvBufCapMax,
proto: proto,
exitChan: make(chan struct{}),
waitGroup: &sync.WaitGroup{},
},
}
}
func (self *TcpClient) Run() {
conn, err := net.DialTCP("tcp", nil, self.svrAddr)
CheckError(err)
self.waitGroup.Add(1)
go func() {
// client sock do NOT need to identify self
c := newTcpConn( /*atomic.AddUint32(&self.autoIncID, 1)*/ 0, self.tcpSock, conn, self.sendBufCap, self.recvBufCap, self.connClose)
if self.onConnConnect != nil {
self.onConnConnect(c)
}
c.run()
self.waitGroup.Done()
}()
}
func (self *TcpClient) Close() {
close(self.exitChan)
self.waitGroup.Wait()
}
func (self *TcpClient) OnConnect(fn OnTcpConnCallback) {
self.onConnConnect = fn
}
func (self *TcpClient) OnClose(fn OnTcpConnCallback) {
self.onConnClose = fn
}
func (self *TcpClient) connClose(conn *TcpConn) {
if self.onConnClose != nil {
self.onConnClose(conn)
}
}
可以看到，tcpsock 的实现中有几处地方使用到了回调函数，或许后面会有更好的封装及处理方法（当然这绝不是说回调函数不好），暂时先这样设计吧。
而我也编写了个简单的包括服务器与客户端的聊天室程序，用以测试及验证 tcpsock，虽然这个示例挺简单，但可能比较实用，包括不限于如何使用 tcpsock，及二进制协议的设计与解析（涉及到 TCP 粘包处理）等，代码如下：
package main
import (
"log"
"os"
"os/signal"
"sync"
"syscall"
"tcpsock"
. "tcpsock/samples/chatroom/protocol"
"time"
. "github.com/ecofast/sysutils"
)
var (
shutdown = make(chan bool, 1)
mutex sync.Mutex
clients map[uint32]*tcpsock.TcpConn
)
func init() {
signals := make(chan os.Signal, 1)
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
go func() {
<-signals
shutdown <- true
}()
}
func onConnConnect(conn *tcpsock.TcpConn) {
// conn.Send(genChatPacket())
mutex.Lock()
defer mutex.Unlock()
clients[conn.ID()] = conn
}
func onConnClose(conn *tcpsock.TcpConn) {
mutex.Lock()
defer mutex.Unlock()
delete(clients, conn.ID())
}
func genChatPacket() *ChatPacket {
var head PacketHead
head.Signature = ChatSignature
head.PlayerID = 555555555
s := "current time is " + TimeToStr(time.Now())
head.BodyLen = uint32(len(s))
body := make([]byte, int(head.BodyLen))
copy(body[:], []byte(s)[:])
return NewChatPacket(head, body)
}
func broadcast() {
mutex.Lock()
defer mutex.Unlock()
packet := genChatPacket()
for _, c := range clients {
c.Write(packet)
}
}
func onMsg(conn *tcpsock.TcpConn, p *ChatPacket) {
mutex.Lock()
defer mutex.Unlock()
for _, c := range clients {
c.Write(p)
}
}
func main() {
clients = make(map[uint32]*tcpsock.TcpConn)
proto := &ChatProtocol{}
proto.OnMessage(onMsg)
server := tcpsock.NewTcpServer(9999, 2, proto)
server.OnConnConnect(onConnConnect)
server.OnConnClose(onConnClose)
log.Println("=====service start=====")
go server.Serve()
ticker := time.NewTicker(10 * time.Second)
go func() {
for range ticker.C {
log.Printf("num of conn: %d\n", server.NumOfConn())
broadcast()
}
}()
<-shutdown
server.Close()
log.Println("=====service end=====")
}
package protocol
import (
"bytes"
"encoding/binary"
"tcpsock"
)
const (
ChatSignature = 0xFFFFFFFF
PacketHeadSize = 4 + 4 + 4
)
type PacketHead struct {
Signature uint32
PlayerID uint32
BodyLen uint32
}
func (head *PacketHead) Bytes() []byte {
var buf bytes.Buffer
binary.Write(&buf, binary.LittleEndian, head)
return buf.Bytes()
}
type ChatPacket struct {
PacketHead
Body []byte
}
func NewChatPacket(head PacketHead, body []byte) *ChatPacket {
return &ChatPacket{
PacketHead: head,
Body: body,
}
}
func (p *ChatPacket) Marshal() []byte {
buf := make([]byte, PacketHeadSize+len(p.Body))
copy(buf[:PacketHeadSize], p.PacketHead.Bytes()[:])
copy(buf[PacketHeadSize:], p.Body[:])
return buf
}
type ChatProtocol struct {
recvBuf []byte
recvBufLen int
onMsg func(c *tcpsock.TcpConn, p *ChatPacket)
}
func (self *ChatProtocol) Parse(b []byte, recvChan chan<- tcpsock.Packet) {
count := len(b)
if count+self.recvBufLen > tcpsock.RecvBufLenMax {
return
}
self.recvBuf = append(self.recvBuf, b[0:count]...)
self.recvBufLen += count
offsize := 0
offset := 0
var head PacketHead
for self.recvBufLen-offsize > PacketHeadSize {
offset = 0
head.Signature = uint32(uint32(self.recvBuf[offsize+3])<<24 | uint32(self.recvBuf[offsize+2])<<16 | uint32(self.recvBuf[offsize+1])<<8 | uint32(self.recvBuf[offsize+0]))
offset += 4
head.PlayerID = uint32(uint32(self.recvBuf[offsize+offset+3])<<24 | uint32(self.recvBuf[offsize+offset+2])<<16 | uint32(self.recvBuf[offsize+offset+1])<<8 | uint32(self.recvBuf[offsize+offset+0]))
offset += 4
head.BodyLen = uint32(uint32(self.recvBuf[offsize+offset+3])<<24 | uint32(self.recvBuf[offsize+offset+2])<<16 | uint32(self.recvBuf[offsize+offset+1])<<8 | uint32(self.recvBuf[offsize+offset+0]))
offset += 4
if head.Signature == ChatSignature {
pkglen := int(PacketHeadSize + head.BodyLen)
if pkglen >= tcpsock.RecvBufLenMax {
offsize = self.recvBufLen
break
}
if offsize+pkglen > self.recvBufLen {
break
}
recvChan <- NewChatPacket(head, self.recvBuf[offsize+offset:offsize+offset+int(head.BodyLen)])
offsize += pkglen
} else {
offsize++
}
}
self.recvBufLen -= offsize
if self.recvBufLen > 0 {
self.recvBuf = self.recvBuf[offsize : offsize+self.recvBufLen]
} else {
self.recvBuf = nil
}
}
func (self *ChatProtocol) Process(conn *tcpsock.TcpConn, p tcpsock.Packet) {
packet := p.(*ChatPacket)
self.onMsg(conn, packet)
}
func (self *ChatProtocol) OnMessage(fn func(c *tcpsock.TcpConn, p *ChatPacket)) {
self.onMsg = fn
}
package main
import (
"fmt"
"log"
"os"
"os/signal"
"syscall"
"tcpsock"
. "tcpsock/samples/chatroom/protocol"
)
const (
ServerAddr = ":9999"
)
var (
shutdown = make(chan bool, 1)
tcpConn *tcpsock.TcpConn
id uint32
)
func init() {
signals := make(chan os.Signal, 1)
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
go func() {
<-signals
shutdown <- true
}()
}
func main() {
genID()
proto := &ChatProtocol{}
proto.OnMessage(onMsg)
client := tcpsock.NewTcpClient(ServerAddr, proto)
client.OnConnect(onConnect)
client.OnClose(onClose)
go client.Run()
go input()
<-shutdown
client.Close()
}
func onConnect(c *tcpsock.TcpConn) {
log.Println("successfully connect to server", c.RawConn().RemoteAddr().String())
tcpConn = c
}
func onClose(c *tcpsock.TcpConn) {
log.Println("disconnect from server", c.RawConn().RemoteAddr().String())
tcpConn = nil
}
func onMsg(c *tcpsock.TcpConn, p *ChatPacket) {
log.Printf("%d: %s\n", p.PlayerID, string(p.Body))
}
func genID() {
fmt.Printf("pls enter your id: ")
fmt.Scan(&id)
fmt.Println("your id is:", id)
}
func input() {
s := ""
for {
if n, err := fmt.Scan(&s); n == 0 || err != nil {
break
}
if tcpConn == nil {
break
}
tcpConn.Write(genPacket(s))
}
}
func genPacket(s string) *ChatPacket {
var head PacketHead
head.Signature = ChatSignature
head.PlayerID = id
head.BodyLen = uint32(len(s))
body := make([]byte, head.BodyLen)
copy(body[:], []byte(s)[:])
return NewChatPacket(head, body)
}
可以看到，服务器与客户端完全共用一套协议处理单元，且各自只需对收到的网络封包作处理即可。
以上是关于tcpsock for Golang的主要内容，如果未能解决你的问题，请参考以下文章
golang 片段7 for https://medium.com/@francesc/why-are-there-nil-channels-in-go-9877cc0b2308