代码优化

wuyingsong · Oct 16, 2017 · afae40a · afae40a
1 parent 13f1592
commit afae40a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 193 deletions.
diff --git a/conn.go b/conn.go
diff --git a/tcp_conn.go b/tcp_conn.go
@@ -2,6 +2,7 @@ package tcp
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"strings"
@@ -26,7 +27,6 @@ type TCPConn struct {
 
 	readDeadline  time.Duration
 	writeDeadline time.Duration
-	maxPacketSize uint32
 
 	exitChan  chan struct{}
 	closeOnce sync.Once
@@ -49,17 +49,23 @@ func NewTCPConn(conn *net.TCPConn, callback CallBack, protocol Protocol) *TCPCon
 	return c
 }
 
-func (c *TCPConn) Serve() {
+func (c *TCPConn) Serve() error {
 	defer func() {
 		if r := recover(); r != nil {
 			logger.Println("tcp conn(%v) Serve error, %v ", c.GetRemoteIPAddress(), r)
 		}
 	}()
+	if c.callback == nil || c.protocol == nil {
+		err := fmt.Errorf("callback and protocol are not allowed to be nil")
+		c.Close()
+		return err
+	}
 	atomic.StoreInt32(&c.exitFlag, 1)
 	c.callback.OnConnected(c)
 	go c.readLoop()
 	go c.writeLoop()
 	go c.handleLoop()
+	return nil
 }
 
 func (c *TCPConn) readLoop() {
@@ -146,7 +152,9 @@ func (c *TCPConn) Close() {
 		close(c.errChan)
 		close(c.writeChan)
 		close(c.readChan)
-		c.callback.OnDisconnected(c)
+		if c.callback != nil {
+			c.callback.OnDisconnected(c)
+		}
 		atomic.StoreInt32(&c.exitFlag, 0)
 		c.conn.Close()
 	})