Mod: writePkgNum count

client_test.go

@@ -126,16 +126,17 @@ func TestTCPClient(t *testing.T) {
 	beforeWriteBytes := atomic.LoadUint32(&conn.writeBytes)
 	beforeWritePkgNum := atomic.LoadUint32(&conn.writePkgNum)
 	_, err = conn.send([]byte("hello"))
+	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&conn.writePkgNum))
 	assert.Nil(t, err)
 	assert.Equal(t, beforeWriteBytes+5, atomic.LoadUint32(&conn.writeBytes))
 	err = ss.WriteBytes([]byte("hello"))
 	assert.Equal(t, beforeWriteBytes+10, atomic.LoadUint32(&conn.writeBytes))
-	assert.Equal(t, beforeWritePkgNum+1, atomic.LoadUint32(&conn.writePkgNum))
+	assert.Equal(t, beforeWritePkgNum+2, atomic.LoadUint32(&conn.writePkgNum))
 	assert.Nil(t, err)
 	var pkgs [][]byte
 	pkgs = append(pkgs, []byte("hello"), []byte("hello"))
 	_, err = conn.send(pkgs)
-	assert.Equal(t, beforeWritePkgNum+3, atomic.LoadUint32(&conn.writePkgNum))
+	assert.Equal(t, beforeWritePkgNum+4, atomic.LoadUint32(&conn.writePkgNum))
 	assert.Equal(t, beforeWriteBytes+20, atomic.LoadUint32(&conn.writeBytes))
 	assert.Nil(t, err)
 	ss.SetCompressType(CompressSnappy)

connection.go

@@ -292,6 +292,7 @@ func (t *gettyTCPConn) send(pkg interface{}) (int, error) {
 	if p, ok = pkg.([]byte); ok {
 		if length, err = t.writer.Write(p); 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)
@@ -449,11 +450,11 @@ func (u *gettyUDPConn) send(udpCtx interface{}) (int, error) {
 
 	if length, _, err = u.conn.WriteMsgUDP(buf, nil, peerAddr); err == nil {
 		atomic.AddUint32(&u.writeBytes, (uint32)(len(buf)))
+		atomic.AddUint32(&u.writePkgNum, 1)
 	}
 	log.Debugf("WriteMsgUDP(peerAddr:%s) = {length:%d, error:%v}", peerAddr, length, err)
 
 	return length, perrors.WithStack(err)
-	//return length, err
 }
 
 // close udp connection
@@ -591,9 +592,9 @@ func (w *gettyWSConn) send(pkg interface{}) (int, error) {
 	w.updateWriteDeadline()
 	if err = w.conn.WriteMessage(websocket.BinaryMessage, p); err == nil {
 		atomic.AddUint32(&w.writeBytes, (uint32)(len(p)))
+		atomic.AddUint32(&w.writePkgNum, 1)
 	}
 	return len(p), perrors.WithStack(err)
-	//return len(p), err
 }
 
 func (w *gettyWSConn) writePing() error {

session.go

@@ -407,7 +407,6 @@ func (s *session) WritePkg(pkg interface{}, timeout time.Duration) error {
 			log.Warnf("%s, [session.WritePkg] @s.Connection.Write(pkg:%#v) = err:%v", s.Stat(), pkg, err)
 			return perrors.WithStack(err)
 		}
-		s.incWritePkgNum()
 		return nil
 	}
 	select {
@@ -432,9 +431,6 @@ func (s *session) WriteBytes(pkg []byte) error {
 	if _, err := s.Connection.send(pkg); err != nil {
 		return perrors.Wrapf(err, "s.Connection.Write(pkg len:%d)", len(pkg))
 	}
-
-	s.incWritePkgNum()
-
 	return nil
 }