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,21 +19,21 @@ 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
 }
 
 //创建连接的方法
@@ -43,7 +45,6 @@ func NewConntion(server ziface.IServer, conn *net.TCPConn, connID uint32, msgHan
 		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{}),
@@ -81,7 +82,7 @@ func (c *Connection) StartWriter() {
 				fmt.Println("msgBuffChan is Closed")
 				break
 			}
-		case <-c.ExitBuffChan:
+		case <-c.ctx.Done():
 			return
 		}
 	}
@@ -96,6 +97,10 @@ func (c *Connection) StartReader() {
 	defer c.Stop()
 
 	for {
+		select {
+		case <-c.ctx.Done():
+			return
+		default:
 			// 创建拆包解包的对象
 			dp := NewDataPack()
 
@@ -139,10 +144,12 @@ func (c *Connection) StartReader() {
 				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)
 }