optimize: add HeaderLen field

cloudwego · Aug 15, 2024 · 29b4982 · 29b4982
1 parent 4a81f2a
commit 29b4982
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/protocol/ttheader/decode.go b/protocol/ttheader/decode.go
@@ -50,6 +50,9 @@ type DecodeParam struct {
 	// You can refer to metakey.go for more details.
 	StrInfo map[string]string
 
+	// HeaderLen is used to set up header length of a request/response.
+	HeaderLen int
+
 	// PayloadLen is used to set up payload length of a request/response.
 	PayloadLen int
 }
@@ -112,7 +115,8 @@ func Decode(ctx context.Context, in bufiox.Reader) (param DecodeParam, err error
 		return
 	}
 
-	param.PayloadLen = int(totalLen - uint32(headerInfoSize) + Size32 - TTHeaderMetaSize)
+	param.HeaderLen = int(uint32(headerInfoSize) + TTHeaderMetaSize)
+	param.PayloadLen = int(totalLen) + Size32 - param.HeaderLen
 	return
 }
 

diff --git a/protocol/ttheader/encode_test.go b/protocol/ttheader/encode_test.go
@@ -54,7 +54,7 @@ func TestEncodeToBytes(t *testing.T) {
 	if err != nil {
 		t.Fatalf("encode to bytes failed, %s", err.Error())
 	}
-	checkParamEqual(t, encodeParam, decodeParam)
+	checkParamEqual(t, encodeParam, decodeParam, len(buf), 0)
 }
 
 func TestEncode(t *testing.T) {
@@ -110,12 +110,13 @@ func TestEncode(t *testing.T) {
 		t.Fatalf("encode failed, %s", err.Error())
 	}
 	binary.BigEndian.PutUint32(totalLenField, uint32(bw.WrittenLen()-4))
+	headerLen := bw.WrittenLen()
 	bw.Flush()
 	wg.Wait()
-	checkParamEqual(t, encodeParam, decodeParam)
+	checkParamEqual(t, encodeParam, decodeParam, headerLen, 0)
 }
 
-func checkParamEqual(t *testing.T, encodeParam EncodeParam, decodeParam DecodeParam) {
+func checkParamEqual(t *testing.T, encodeParam EncodeParam, decodeParam DecodeParam, headerLen, payloadLen int) {
 	if decodeParam.Flags != encodeParam.Flags {
 		t.Fatalf("encode to bytes failed, flags not equal")
 	}
@@ -131,7 +132,10 @@ func checkParamEqual(t *testing.T, encodeParam EncodeParam, decodeParam DecodePa
 	if !reflect.DeepEqual(decodeParam.StrInfo, encodeParam.StrInfo) {
 		t.Fatalf("encode to bytes failed, str info not equal")
 	}
-	if decodeParam.PayloadLen != 0 {
+	if decodeParam.HeaderLen != headerLen {
+		t.Fatalf("encode to bytes failed, header len not equal")
+	}
+	if decodeParam.PayloadLen != payloadLen {
 		t.Fatalf("encode to bytes failed, payload len not equal")
 	}
 }