From 7352d7b91c75d3a453f15c4140cf87877a65c155 Mon Sep 17 00:00:00 2001
From: iosmanthus <MyOsmanthusTree@gmail.com>
Date: Sun, 12 Feb 2023 00:04:02 +0800
Subject: [PATCH] *: reduce the allocation of error constructing in
 `DecodeTableID` (#41301)

close pingcap/tidb#41300
---
 tablecodec/BUILD.bazel        |  1 +
 tablecodec/tablecodec.go      | 16 ++++++++++------
 tablecodec/tablecodec_test.go | 19 +++++++++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/tablecodec/BUILD.bazel b/tablecodec/BUILD.bazel
index db89b945b183d..2cf712bd02d34 100644
--- a/tablecodec/BUILD.bazel
+++ b/tablecodec/BUILD.bazel
@@ -49,6 +49,7 @@ go_test(
         "//util/rowcodec",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
+        "@com_github_tikv_client_go_v2//tikv",
         "@org_uber_go_goleak//:goleak",
     ],
 )
diff --git a/tablecodec/tablecodec.go b/tablecodec/tablecodec.go
index 394277ba2768c..e5adfa6f91b32 100644
--- a/tablecodec/tablecodec.go
+++ b/tablecodec/tablecodec.go
@@ -278,13 +278,17 @@ func DecodeKeyHead(key kv.Key) (tableID int64, indexID int64, isRecordKey bool,
 
 // DecodeTableID decodes the table ID of the key, if the key is not table key, returns 0.
 func DecodeTableID(key kv.Key) int64 {
-	// If the key is in API V2, then ignore the prefix
-	_, k, err := tikv.DecodeKey(key, kvrpcpb.APIVersion_V2)
-	if err == nil {
-		key = k
-	}
 	if !key.HasPrefix(tablePrefix) {
-		return 0
+		// If the key is in API V2, then ignore the prefix
+		_, k, err := tikv.DecodeKey(key, kvrpcpb.APIVersion_V2)
+		if err != nil {
+			terror.Log(errors.Trace(err))
+			return 0
+		}
+		key = k
+		if !key.HasPrefix(tablePrefix) {
+			return 0
+		}
 	}
 	key = key[len(tablePrefix):]
 	_, tableID, err := codec.DecodeInt(key)
diff --git a/tablecodec/tablecodec_test.go b/tablecodec/tablecodec_test.go
index adc4ccc78c13b..45020e034938c 100644
--- a/tablecodec/tablecodec_test.go
+++ b/tablecodec/tablecodec_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/rowcodec"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/tikv"
 )
 
 // TestTableCodec  tests some functions in package tablecodec
@@ -698,3 +699,21 @@ func TestTempIndexValueCodec(t *testing.T) {
 	require.Equal(t, result[1].Handle.IntValue(), int64(100))
 	require.Equal(t, result[2].Handle.IntValue(), int64(101))
 }
+
+func TestV2TableCodec(t *testing.T) {
+	const tableID int64 = 31415926
+	key := EncodeTablePrefix(tableID)
+	c, err := tikv.NewCodecV2(tikv.ModeTxn, 271828)
+	require.NoError(t, err)
+	key = c.EncodeKey(key)
+	tbid := DecodeTableID(key)
+	require.Equal(t, tableID, tbid)
+
+	key = []byte("x001HelloWorld")
+	tbid = DecodeTableID(key)
+	require.Equal(t, int64(0), tbid)
+
+	key = []byte("x001x001t123")
+	tbid = DecodeTableID(key)
+	require.Equal(t, int64(0), tbid)
+}