Merge branch 'master' into improve/delte-handleLoop

client_test.go

@@ -135,27 +135,31 @@ func TestTCPClient(t *testing.T) {
 	ss.SetCompressType(CompressNone)
 	conn := ss.(*session).Connection.(*gettyTCPConn)
 	assert.True(t, conn.compress == CompressNone)
-	beforeWriteBytes := conn.writeBytes.Load()
-	beforeWritePkgNum := conn.writePkgNum.Load()
-	_, err = conn.send([]byte("hello"))
-	assert.Equal(t, beforeWritePkgNum+1, conn.writePkgNum.Load())
+	beforeWriteBytes := atomic.LoadUint32(&conn.writeBytes)
+	beforeWritePkgNum := atomic.LoadUint32(&conn.writePkgNum)
+	l, err := conn.send([]byte("hello"))
 	assert.Nil(t, err)
-	assert.Equal(t, beforeWriteBytes+5, conn.writeBytes.Load())
-	err = ss.WriteBytes([]byte("hello"))
-	assert.Equal(t, beforeWriteBytes+10, conn.writeBytes.Load())
-	assert.Equal(t, beforeWritePkgNum+2, conn.writePkgNum.Load())
+	assert.True(t, l == 5)
+	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&conn.writePkgNum))
+	assert.Equal(t, beforeWriteBytes+5, atomic.LoadUint32(&conn.writeBytes))
+	l, err = ss.WriteBytes([]byte("hello"))
 	assert.Nil(t, err)
+	assert.True(t, l == 5)
+	assert.Equal(t, beforeWriteBytes+10, atomic.LoadUint32(&conn.writeBytes))
+	assert.Equal(t, beforeWritePkgNum+2, atomic.LoadUint32(&conn.writePkgNum))
 	var pkgs [][]byte
 	pkgs = append(pkgs, []byte("hello"), []byte("hello"))
-	_, err = conn.send(pkgs)
-	assert.Equal(t, beforeWritePkgNum+4, conn.writePkgNum.Load())
-	assert.Equal(t, beforeWriteBytes+20, conn.writeBytes.Load())
+	l, err = conn.send(pkgs)
 	assert.Nil(t, err)
+	assert.True(t, l == 10)
+	assert.Equal(t, beforeWritePkgNum+4, atomic.LoadUint32(&conn.writePkgNum))
+	assert.Equal(t, beforeWriteBytes+20, atomic.LoadUint32(&conn.writeBytes))
 	ss.SetCompressType(CompressSnappy)
-	err = ss.WriteBytesArray(pkgs...)
+	l, err = ss.WriteBytesArray(pkgs...)
 	assert.Nil(t, err)
-	assert.Equal(t, beforeWritePkgNum+6, conn.writePkgNum.Load())
-	assert.Equal(t, beforeWriteBytes+30, conn.writeBytes.Load())
+	assert.True(t, l == 10)
+	assert.Equal(t, beforeWritePkgNum+6, atomic.LoadUint32(&conn.writePkgNum))
+	assert.Equal(t, beforeWriteBytes+30, atomic.LoadUint32(&conn.writeBytes))
 	assert.True(t, conn.compress == CompressSnappy)
 
 	clt.Close()
@@ -166,6 +170,8 @@ func TestUDPClient(t *testing.T) {
 	var (
 		err      error
 		conn     *net.UDPConn
+		sendLen  int
+		totalLen int
 	)
 	func() {
 		ip := net.ParseIP("127.0.0.1")
@@ -199,15 +205,22 @@ func TestUDPClient(t *testing.T) {
 
 	assert.Equal(t, 1, msgHandler.SessionNumber())
 	ss := msgHandler.array[0]
-	err = ss.WritePkg(nil, 0)
+	totalLen, sendLen, err = ss.WritePkg(nil, 0)
 	assert.NotNil(t, err)
-	err = ss.WritePkg([]byte("hello"), 0)
+	assert.True(t, sendLen == 0)
+	assert.True(t, totalLen == 0)
+	totalLen, sendLen, err = ss.WritePkg([]byte("hello"), 0)
 	assert.NotNil(t, perrors.Cause(err))
-	err = ss.WriteBytes([]byte("hello"))
+	assert.True(t, sendLen == 0)
+	assert.True(t, totalLen == 0)
+	l, err := ss.WriteBytes([]byte("hello"))
+	assert.Zero(t, l)
 	assert.NotNil(t, err)
-	err = ss.WriteBytesArray([]byte("hello"))
+	l, err = ss.WriteBytesArray([]byte("hello"))
+	assert.Zero(t, l)
 	assert.NotNil(t, err)
-	err = ss.WriteBytesArray([]byte("hello"), []byte("world"))
+	l, err = ss.WriteBytesArray([]byte("hello"), []byte("world"))
+	assert.Zero(t, l)
 	assert.NotNil(t, err)
 	ss.SetCompressType(CompressNone)
 	host, port, _ := net.SplitHostPort(addr.String())
@@ -230,10 +243,12 @@ func TestUDPClient(t *testing.T) {
 	assert.Equal(t, beforeWriteBytes+5, udpConn.writeBytes.Load())
 	assert.Nil(t, err)
 
-	beforeWritePkgNum := udpConn.writePkgNum.Load()
-	err = ss.WritePkg(udpCtx, 0)
-	assert.Equal(t, beforeWritePkgNum+1, udpConn.writePkgNum.Load())
+	beforeWritePkgNum := atomic.LoadUint32(&udpConn.writePkgNum)
+	totalLen, sendLen, err = ss.WritePkg(udpCtx, 0)
+	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&udpConn.writePkgNum))
 	assert.Nil(t, err)
+	assert.True(t, sendLen == 0)
+	assert.True(t, totalLen == 0)
 
 	clt.Close()
 	assert.True(t, clt.IsClosed())
@@ -283,19 +298,22 @@ func TestNewWSClient(t *testing.T) {
 	assert.True(t, conn.compress == CompressNone)
 	err := conn.handlePing("hello")
 	assert.Nil(t, err)
-	_, err = conn.send("hello")
+	l, err := conn.send("hello")
 	assert.NotNil(t, err)
-	beforeWriteBytes := conn.writeBytes.Load()
+	assert.True(t, l == 0)
+	beforeWriteBytes := atomic.LoadUint32(&conn.writeBytes)
 	_, err = conn.send([]byte("hello"))
 	assert.Nil(t, err)
-	assert.Equal(t, beforeWriteBytes+5, conn.writeBytes.Load())
-	beforeWritePkgNum := conn.writePkgNum.Load()
-	err = ss.WriteBytes([]byte("hello"))
+	assert.Equal(t, beforeWriteBytes+5, atomic.LoadUint32(&conn.writeBytes))
+	beforeWritePkgNum := atomic.LoadUint32(&conn.writePkgNum)
+	l, err = ss.WriteBytes([]byte("hello"))
 	assert.Nil(t, err)
-	assert.Equal(t, beforeWritePkgNum+1, conn.writePkgNum.Load())
-	err = ss.WriteBytesArray([]byte("hello"), []byte("hello"))
+	assert.True(t, l == 5)
+	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&conn.writePkgNum))
+	l, err = ss.WriteBytesArray([]byte("hello"), []byte("hello"))
 	assert.Nil(t, err)
-	assert.Equal(t, beforeWritePkgNum+3, conn.writePkgNum.Load())
+	assert.True(t, l == 10)
+	assert.Equal(t, beforeWritePkgNum+3, atomic.LoadUint32(&conn.writePkgNum))
 	err = conn.writePing()
 	assert.Nil(t, err)
 

connection.go

@@ -268,6 +268,7 @@ func (t *gettyTCPConn) send(pkg interface{}) (int, error) {
 		ok          bool
 		p           []byte
 		length      int
+		lg          int64
 	)
 
 	if t.compress == CompressNone && t.wTimeout > 0 {
@@ -285,19 +286,21 @@ func (t *gettyTCPConn) send(pkg interface{}) (int, error) {
 
 	if buffers, ok := pkg.([][]byte); ok {
 		netBuf := net.Buffers(buffers)
-		if length, err := netBuf.WriteTo(t.conn); err == nil {
-			t.writeBytes.Add((uint32)(length))
-			t.writePkgNum.Add((uint32)(len(buffers)))
+		lg, err = netBuf.WriteTo(t.conn)
+		if err == nil {
+			atomic.AddUint32(&t.writeBytes, (uint32)(lg))
+			atomic.AddUint32(&t.writePkgNum, (uint32)(len(buffers)))
 		}
 		log.Debugf("localAddr: %s, remoteAddr:%s, now:%s, length:%d, err:%s",
 			t.conn.LocalAddr(), t.conn.RemoteAddr(), currentTime, length, err)
-		return int(length), perrors.WithStack(err)
+		return int(lg), perrors.WithStack(err)
 	}
 
 	if p, ok = pkg.([]byte); ok {
-		if length, err = t.writer.Write(p); err == nil {
-			t.writeBytes.Add((uint32)(len(p)))
-			t.writePkgNum.Add(1)
+		length, err = t.writer.Write(p)
+		if err == nil {
+			atomic.AddUint32(&t.writeBytes, (uint32)(len(p)))
+			atomic.AddUint32(&t.writePkgNum, 1)
 		}
 		log.Debugf("localAddr: %s, remoteAddr:%s, now:%s, length:%d, err:%s",
 			t.conn.LocalAddr(), t.conn.RemoteAddr(), currentTime, length, err)

demo/hello/hello.go

@@ -31,7 +31,7 @@ func ClientRequest() {
 		go func() {
 			echoTimes := 10
 			for i := 0; i < echoTimes; i++ {
-				err := ss.WritePkg("hello", WritePkgTimeout)
+				_, _, err := ss.WritePkg("hello", WritePkgTimeout)
 				if err != nil {
 					log.Infof("session.WritePkg(session{%s}, error{%v}", ss.Stat(), err)
 					ss.Close()

getty.go

@@ -171,9 +171,12 @@ type Session interface {
 
 	// the Writer will invoke this function. Pls attention that if timeout is less than 0, WritePkg will send @pkg asap.
 	// for udp session, the first parameter should be UDPContext.
-	WritePkg(pkg interface{}, timeout time.Duration) error
-	WriteBytes([]byte) error
-	WriteBytesArray(...[]byte) error
+	// totalBytesLength: @pkg stream bytes length after encoding @pkg.
+	// sendBytesLength: stream bytes length that sent out successfully.
+	// err: maybe it has illegal data, encoding error, or write out system error.
+	WritePkg(pkg interface{}, timeout time.Duration) (totalBytesLength int, sendBytesLength int, err error)
+	WriteBytes([]byte) (int, error)
+	WriteBytesArray(...[]byte) (int, error)
 	Close()
 }
 

server_test.go

@@ -76,6 +76,7 @@ func testTCPServer(t *testing.T, address string) {
 	server.Close()
 	assert.True(t, server.IsClosed())
 }
+
 func testTCPTlsServer(t *testing.T, address string) {
 	var (
 		server           *server

session.go

@@ -338,12 +338,12 @@ func (s *session) sessionToken() string {
 		s.name, s.EndPoint().EndPointType(), s.ID(), s.LocalAddr(), s.RemoteAddr())
 }
 
-func (s *session) WritePkg(pkg interface{}, timeout time.Duration) error {
+func (s *session) WritePkg(pkg interface{}, timeout time.Duration) (int, int, error) {
 	if pkg == nil {
-		return fmt.Errorf("@pkg is nil")
+		return 0, 0, fmt.Errorf("@pkg is nil")
 	}
 	if s.IsClosed() {
-		return ErrSessionClosed
+		return 0, 0, ErrSessionClosed
 	}
 
 	defer func() {
@@ -358,7 +358,7 @@ func (s *session) WritePkg(pkg interface{}, timeout time.Duration) error {
 	pkgBytes, err := s.writer.Write(s, pkg)
 	if err != nil {
 		log.Warnf("%s, [session.WritePkg] session.writer.Write(@pkg:%#v) = error:%+v", s.Stat(), pkg, err)
-		return perrors.WithStack(err)
+		return len(pkgBytes), 0, perrors.WithStack(err)
 	}
 	var udpCtxPtr *UDPContext
 	if udpCtx, ok := pkg.(UDPContext); ok {
@@ -375,32 +375,32 @@ func (s *session) WritePkg(pkg interface{}, timeout time.Duration) error {
 	if 0 < timeout {
 		s.Connection.SetWriteTimeout(timeout)
 	}
-	_, err = s.Connection.send(pkg)
+	var succssCount int
+	succssCount, err = s.Connection.send(pkg)
 	if err != nil {
 		log.Warnf("%s, [session.WritePkg] @s.Connection.Write(pkg:%#v) = err:%+v", s.Stat(), pkg, err)
-		return perrors.WithStack(err)
+		return len(pkgBytes), succssCount, perrors.WithStack(err)
 	}
-
-	return nil
+	return len(pkgBytes), succssCount, nil
 }
 
 // for codecs
-func (s *session) WriteBytes(pkg []byte) error {
+func (s *session) WriteBytes(pkg []byte) (int, error) {
 	if s.IsClosed() {
-		return ErrSessionClosed
+		return 0, ErrSessionClosed
 	}
 
-	// s.conn.SetWriteTimeout(time.Now().Add(s.wTimeout))
-	if _, err := s.Connection.send(pkg); err != nil {
-		return perrors.Wrapf(err, "s.Connection.Write(pkg len:%d)", len(pkg))
+	lg, err := s.Connection.send(pkg)
+	if err != nil {
+		return 0, perrors.Wrapf(err, "s.Connection.Write(pkg len:%d)", len(pkg))
 	}
-	return nil
+	return lg, nil
 }
 
 // Write multiple packages at once. so we invoke write sys.call just one time.
-func (s *session) WriteBytesArray(pkgs ...[]byte) error {
+func (s *session) WriteBytesArray(pkgs ...[]byte) (int, error) {
 	if s.IsClosed() {
-		return ErrSessionClosed
+		return 0, ErrSessionClosed
 	}
 	if len(pkgs) == 1 {
 		return s.WriteBytes(pkgs[0])
@@ -408,15 +408,17 @@ func (s *session) WriteBytesArray(pkgs ...[]byte) error {
 
 	// reduce syscall and memcopy for multiple packages
 	if _, ok := s.Connection.(*gettyTCPConn); ok {
-		if _, err := s.Connection.send(pkgs); err != nil {
-			return perrors.Wrapf(err, "s.Connection.Write(pkgs num:%d)", len(pkgs))
+		lg, err := s.Connection.send(pkgs)
+		if err != nil {
+			return 0, perrors.Wrapf(err, "s.Connection.Write(pkgs num:%d)", len(pkgs))
 		}
-		return nil
+		return lg, nil
 	}
 
 	// get len
 	var (
 		l      int
+		wlg    int
 		err    error
 		length int
 		arrp   *[]byte
@@ -428,7 +430,6 @@ func (s *session) WriteBytesArray(pkgs ...[]byte) error {
 	}
 
 	// merge the pkgs
-	//arr = make([]byte, length)
 	arrp = gxbytes.AcquireBytes(length)
 	defer gxbytes.ReleaseBytes(arrp)
 	arr = *arrp
@@ -439,8 +440,9 @@ func (s *session) WriteBytesArray(pkgs ...[]byte) error {
 		l += len(pkgs[i])
 	}
 
-	if err = s.WriteBytes(arr); err != nil {
-		return perrors.WithStack(err)
+	wlg, err = s.WriteBytes(arr)
+	if err != nil {
+		return 0, perrors.WithStack(err)
 	}
 
 	num := len(pkgs) - 1
@@ -448,12 +450,12 @@ func (s *session) WriteBytesArray(pkgs ...[]byte) error {
 		s.incWritePkgNum()
 	}
 
-	return nil
+	return wlg, nil
 }
 
 func sessionTimerLoop(_ gxtime.TimerID, _ time.Time, arg interface{}) error {
 	ss, _ := arg.(*session)
-	if ss != nil && ss.IsClosed() {
+	if ss == nil || ss.IsClosed() {
 		return ErrSessionClosed
 	}