Optimize readmsg

CESSProject · Jan 4, 2023 · a878a14 · a878a14
1 parent 73ecb97
commit a878a14
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/node/transfer.go b/node/transfer.go
@@ -82,9 +82,10 @@ func (t *TcpCon) sendMsg() {
 
 func (t *TcpCon) readMsg() {
 	var (
-		err    error
-		n      int
-		header = make([]byte, 4)
+		err      error
+		n        int
+		waittime int64
+		header   = make([]byte, 4)
 	)
 	readBuf := readBufPool.Get().([]byte)
 	defer func() {
@@ -93,15 +94,22 @@ func (t *TcpCon) readMsg() {
 		close(t.recv)
 		readBufPool.Put(readBuf)
 	}()
-	for {
+	for !t.IsClose() {
 		// read until we get 4 bytes for the magic
 		_, err = io.ReadFull(t.conn, header)
 		if err != nil {
 			if err != io.EOF {
 				err = fmt.Errorf("initial read error: %v \n", err)
 				return
 			}
-			continue
+			if err == io.EOF {
+				waittime++
+				if waittime >= 10 {
+					return
+				}
+				time.Sleep(time.Second)
+				continue
+			}
 		}
 
 		if !bytes.Equal(header, HEAD_FILE) && !bytes.Equal(header, HEAD_FILLER) {