Merge pull request #77 from Zereker/master

MOD: 使用 context 优化 close chan, 读写的时候 conn closed 判断非线程安全

znet/connection.go

@@ -1,13 +1,15 @@
 package znet
 
 import (
+	"context"
 	"errors"
 	"fmt"
-	"github.com/aceld/zinx/utils"
-	"github.com/aceld/zinx/ziface"
 	"io"
 	"net"
 	"sync"
+
+	"github.com/aceld/zinx/utils"
+	"github.com/aceld/zinx/ziface"
 )
 
 type Connection struct {
@@ -17,36 +19,35 @@ type Connection struct {
 	Conn *net.TCPConn
 	//当前连接的ID 也可以称作为SessionID，ID全局唯一
 	ConnID uint32
-	//当前连接的关闭状态
-	isClosed bool
 	//消息管理MsgId和对应处理方法的消息管理模块
 	MsgHandler ziface.IMsgHandle
 	//告知该链接已经退出/停止的channel
-	ExitBuffChan chan bool
+	ctx    context.Context
+	cancel context.CancelFunc
 	//无缓冲管道，用于读、写两个goroutine之间的消息通信
 	msgChan chan []byte
 	//有缓冲管道，用于读、写两个goroutine之间的消息通信
 	msgBuffChan chan []byte
 
+	sync.RWMutex
 	//链接属性
 	property map[string]interface{}
-	//保护链接属性修改的锁
-	propertyLock sync.RWMutex
+	//当前连接的关闭状态
+	isClosed bool
 }
 
 //创建连接的方法
 func NewConntion(server ziface.IServer, conn *net.TCPConn, connID uint32, msgHandler ziface.IMsgHandle) *Connection {
 	//初始化Conn属性
 	c := &Connection{
-		TcpServer:    server,
-		Conn:         conn,
-		ConnID:       connID,
-		isClosed:     false,
-		MsgHandler:   msgHandler,
-		ExitBuffChan: make(chan bool, 1),
-		msgChan:      make(chan []byte),
-		msgBuffChan:  make(chan []byte, utils.GlobalObject.MaxMsgChanLen),
-		property:     make(map[string]interface{}),
+		TcpServer:   server,
+		Conn:        conn,
+		ConnID:      connID,
+		isClosed:    false,
+		MsgHandler:  msgHandler,
+		msgChan:     make(chan []byte),
+		msgBuffChan: make(chan []byte, utils.GlobalObject.MaxMsgChanLen),
+		property:    make(map[string]interface{}),
 	}
 
 	//将新创建的Conn添加到链接管理中
@@ -81,7 +82,7 @@ func (c *Connection) StartWriter() {
 				fmt.Println("msgBuffChan is Closed")
 				break
 			}
-		case <-c.ExitBuffChan:
+		case <-c.ctx.Done():
 			return
 		}
 	}
@@ -96,53 +97,59 @@ func (c *Connection) StartReader() {
 	defer c.Stop()
 
 	for {
-		// 创建拆包解包的对象
-		dp := NewDataPack()
-
-		//读取客户端的Msg head
-		headData := make([]byte, dp.GetHeadLen())
-		if _, err := io.ReadFull(c.Conn, headData); err != nil {
-			fmt.Println("read msg head error ", err)
-			break
-		}
-		//fmt.Printf("read headData %+v\n", headData)
-
-		//拆包，得到msgid 和 datalen 放在msg中
-		msg, err := dp.Unpack(headData)
-		if err != nil {
-			fmt.Println("unpack error ", err)
-			break
-		}
+		select {
+		case <-c.ctx.Done():
+			return
+		default:
+			// 创建拆包解包的对象
+			dp := NewDataPack()
+
+			//读取客户端的Msg head
+			headData := make([]byte, dp.GetHeadLen())
+			if _, err := io.ReadFull(c.Conn, headData); err != nil {
+				fmt.Println("read msg head error ", err)
+				break
+			}
+			//fmt.Printf("read headData %+v\n", headData)
 
-		//根据 dataLen 读取 data，放在msg.Data中
-		var data []byte
-		if msg.GetDataLen() > 0 {
-			data = make([]byte, msg.GetDataLen())
-			if _, err := io.ReadFull(c.Conn, data); err != nil {
-				fmt.Println("read msg data error ", err)
+			//拆包，得到msgid 和 datalen 放在msg中
+			msg, err := dp.Unpack(headData)
+			if err != nil {
+				fmt.Println("unpack error ", err)
 				break
 			}
-		}
-		msg.SetData(data)
 
-		//得到当前客户端请求的Request数据
-		req := Request{
-			conn: c,
-			msg:  msg,
-		}
+			//根据 dataLen 读取 data，放在msg.Data中
+			var data []byte
+			if msg.GetDataLen() > 0 {
+				data = make([]byte, msg.GetDataLen())
+				if _, err := io.ReadFull(c.Conn, data); err != nil {
+					fmt.Println("read msg data error ", err)
+					break
+				}
+			}
+			msg.SetData(data)
 
-		if utils.GlobalObject.WorkerPoolSize > 0 {
-			//已经启动工作池机制，将消息交给Worker处理
-			c.MsgHandler.SendMsgToTaskQueue(&req)
-		} else {
-			//从绑定好的消息和对应的处理方法中执行对应的Handle方法
-			go c.MsgHandler.DoMsgHandler(&req)
+			//得到当前客户端请求的Request数据
+			req := Request{
+				conn: c,
+				msg:  msg,
+			}
+
+			if utils.GlobalObject.WorkerPoolSize > 0 {
+				//已经启动工作池机制，将消息交给Worker处理
+				c.MsgHandler.SendMsgToTaskQueue(&req)
+			} else {
+				//从绑定好的消息和对应的处理方法中执行对应的Handle方法
+				go c.MsgHandler.DoMsgHandler(&req)
+			}
 		}
 	}
 }
 
 //启动连接，让当前连接开始工作
 func (c *Connection) Start() {
+	c.ctx, c.cancel = context.WithCancel(context.Background())
 	//1 开启用户从客户端读取数据流程的Goroutine
 	go c.StartReader()
 	//2 开启用于写回客户端数据流程的Goroutine
@@ -155,6 +162,9 @@ func (c *Connection) Start() {
 func (c *Connection) Stop() {
 	fmt.Println("Conn Stop()...ConnID = ", c.ConnID)
 	//如果当前链接已经关闭
+	c.Lock()
+	defer c.Unlock()
+
 	if c.isClosed == true {
 		return
 	}
@@ -166,13 +176,12 @@ func (c *Connection) Stop() {
 	// 关闭socket链接
 	c.Conn.Close()
 	//关闭Writer
-	c.ExitBuffChan <- true
+	c.cancel()
 
 	//将链接从连接管理器中删除
 	c.TcpServer.GetConnMgr().Remove(c)
 
 	//关闭该链接全部管道
-	close(c.ExitBuffChan)
 	close(c.msgBuffChan)
 }
 
@@ -193,9 +202,13 @@ func (c *Connection) RemoteAddr() net.Addr {
 
 //直接将Message数据发送数据给远程的TCP客户端
 func (c *Connection) SendMsg(msgId uint32, data []byte) error {
+	c.RLock()
 	if c.isClosed == true {
+		c.RUnlock()
 		return errors.New("connection closed when send msg")
 	}
+	c.RUnlock()
+
 	//将data封包，并且发送
 	dp := NewDataPack()
 	msg, err := dp.Pack(NewMsgPackage(msgId, data))
@@ -211,9 +224,13 @@ func (c *Connection) SendMsg(msgId uint32, data []byte) error {
 }
 
 func (c *Connection) SendBuffMsg(msgId uint32, data []byte) error {
+	c.RLock()
 	if c.isClosed == true {
+		c.RUnlock()
 		return errors.New("Connection closed when send buff msg")
 	}
+	c.RUnlock()
+
 	//将data封包，并且发送
 	dp := NewDataPack()
 	msg, err := dp.Pack(NewMsgPackage(msgId, data))
@@ -230,16 +247,16 @@ func (c *Connection) SendBuffMsg(msgId uint32, data []byte) error {
 
 //设置链接属性
 func (c *Connection) SetProperty(key string, value interface{}) {
-	c.propertyLock.Lock()
-	defer c.propertyLock.Unlock()
+	c.Lock()
+	defer c.Unlock()
 
 	c.property[key] = value
 }
 
 //获取链接属性
 func (c *Connection) GetProperty(key string) (interface{}, error) {
-	c.propertyLock.RLock()
-	defer c.propertyLock.RUnlock()
+	c.RLock()
+	defer c.RUnlock()
 
 	if value, ok := c.property[key]; ok {
 		return value, nil
@@ -250,8 +267,8 @@ func (c *Connection) GetProperty(key string) (interface{}, error) {
 
 //移除链接属性
 func (c *Connection) RemoveProperty(key string) {
-	c.propertyLock.Lock()
-	defer c.propertyLock.Unlock()
+	c.Lock()
+	defer c.Unlock()
 
 	delete(c.property, key)
 }