support return total need transfer length

client_test.go

@@ -170,9 +170,10 @@ func TestTCPClient(t *testing.T) {
 
 func TestUDPClient(t *testing.T) {
 	var (
-		err  error
-		conn *net.UDPConn
-		leng int
+		err      error
+		conn     *net.UDPConn
+		leng     int
+		totalLen int
 	)
 	func() {
 		ip := net.ParseIP("127.0.0.1")
@@ -206,12 +207,14 @@ func TestUDPClient(t *testing.T) {
 
 	assert.Equal(t, 1, msgHandler.SessionNumber())
 	ss := msgHandler.array[0]
-	leng, err = ss.WritePkg(nil, 0)
+	totalLen, leng, err = ss.WritePkg(nil, 0)
 	assert.NotNil(t, err)
 	assert.True(t, leng == 0)
-	leng, err = ss.WritePkg([]byte("hello"), 0)
+	assert.True(t, totalLen == 0)
+	totalLen, leng, err = ss.WritePkg([]byte("hello"), 0)
 	assert.NotNil(t, perrors.Cause(err))
 	assert.True(t, leng == 0)
+	assert.True(t, totalLen == 0)
 	l, err := ss.WriteBytes([]byte("hello"))
 	assert.Zero(t, l)
 	assert.NotNil(t, err)
@@ -243,10 +246,11 @@ func TestUDPClient(t *testing.T) {
 	assert.Nil(t, err)
 
 	beforeWritePkgNum := atomic.LoadUint32(&udpConn.writePkgNum)
-	leng, err = ss.WritePkg(udpCtx, 0)
+	totalLen, leng, err = ss.WritePkg(udpCtx, 0)
 	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&udpConn.writePkgNum))
 	assert.Nil(t, err)
 	assert.True(t, leng == 0)
+	assert.True(t, totalLen == 0)
 
 	clt.Close()
 	assert.True(t, clt.IsClosed())

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,7 +171,9 @@ 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) (int, error)
+	// total: need send total bytes
+	// successCount: success send bytes,
+	WritePkg(pkg interface{}, timeout time.Duration) (total int, successCount int, e error)
 	WriteBytes([]byte) (int, error)
 	WriteBytesArray(...[]byte) (int, error)
 	Close()

session.go

@@ -347,12 +347,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) (int, error) {
+func (s *session) WritePkg(pkg interface{}, timeout time.Duration) (int, int, error) {
 	if pkg == nil {
-		return 0, fmt.Errorf("@pkg is nil")
+		return 0, 0, fmt.Errorf("@pkg is nil")
 	}
 	if s.IsClosed() {
-		return 0, ErrSessionClosed
+		return 0, 0, ErrSessionClosed
 	}
 
 	defer func() {
@@ -367,7 +367,7 @@ func (s *session) WritePkg(pkg interface{}, timeout time.Duration) (int, 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 0, perrors.WithStack(err)
+		return len(pkgBytes), 0, perrors.WithStack(err)
 	}
 	var udpCtxPtr *UDPContext
 	if udpCtx, ok := pkg.(UDPContext); ok {
@@ -384,13 +384,13 @@ func (s *session) WritePkg(pkg interface{}, timeout time.Duration) (int, error)
 	if 0 < timeout {
 		s.Connection.SetWriteTimeout(timeout)
 	}
-	var len int
-	len, 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 len, perrors.WithStack(err)
+		return len(pkgBytes), succssCount, perrors.WithStack(err)
 	}
-	return len, nil
+	return len(pkgBytes), succssCount, nil
 }
 
 // for codecs