From 56bc32daad19b9dff10104c55300292de959fde3 Mon Sep 17 00:00:00 2001
From: glorv <glorvs@163.com>
Date: Mon, 14 Dec 2020 15:09:12 +0800
Subject: [PATCH] backend/local: fix next key (#523)

* fix next key

* fix integration test
---
 lightning/backend/local.go                    | 12 +++++-
 lightning/backend/local_test.go               | 40 +++++++++++++++++++
 .../data/vt.precise_types-schema.sql          |  2 +-
 tests/various_types/run.sh                    |  2 +-
 4 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/lightning/backend/local.go b/lightning/backend/local.go
index 33183a12d..1ecf6dd46 100644
--- a/lightning/backend/local.go
+++ b/lightning/backend/local.go
@@ -41,8 +41,8 @@ import (
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/parser/model"
-	"github.com/pingcap/tidb-lightning/lightning/glue"
 	"github.com/pingcap/tidb/table"
+	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/hack"
 	pd "github.com/tikv/pd/client"
@@ -53,6 +53,7 @@ import (
 	"google.golang.org/grpc/keepalive"
 
 	"github.com/pingcap/tidb-lightning/lightning/common"
+	"github.com/pingcap/tidb-lightning/lightning/glue"
 	"github.com/pingcap/tidb-lightning/lightning/log"
 	"github.com/pingcap/tidb-lightning/lightning/manual"
 	"github.com/pingcap/tidb-lightning/lightning/worker"
@@ -1330,6 +1331,15 @@ func nextKey(key []byte) []byte {
 	if len(key) == 0 {
 		return []byte{}
 	}
+
+	// in tikv <= 4.x, tikv will truncate the row key, so we should fetch the next valid row key
+	// See: https://github.com/tikv/tikv/blob/f7f22f70e1585d7ca38a59ea30e774949160c3e8/components/raftstore/src/coprocessor/split_observer.rs#L36-L41
+	if tablecodec.IsRecordKey(key) {
+		tableId, handle, _ := tablecodec.DecodeRecordKey(key)
+		return tablecodec.EncodeRowKeyWithHandle(tableId, handle.Next())
+	}
+
+	// if key is an index, directly append a 0x00 to the key.
 	res := make([]byte, 0, len(key)+1)
 	res = append(res, key...)
 	res = append(res, 0)
diff --git a/lightning/backend/local_test.go b/lightning/backend/local_test.go
index 25c7b28fa..ab55cf113 100644
--- a/lightning/backend/local_test.go
+++ b/lightning/backend/local_test.go
@@ -22,6 +22,11 @@ import (
 
 	"github.com/cockroachdb/pebble"
 	. "github.com/pingcap/check"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/sessionctx/stmtctx"
+	"github.com/pingcap/tidb/tablecodec"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/hack"
 )
 
@@ -49,6 +54,41 @@ func (s *localSuite) TestNextKey(c *C) {
 	// another test case, nextkey()'s return should be smaller than key with a prefix of the origin key
 	next = nextKey([]byte{1, 255})
 	c.Assert(bytes.Compare(next, []byte{1, 255, 0, 1, 2}), Equals, -1)
+
+	// test recode key
+	// key with int handle
+	for _, handleId := range []int64{1, 255, math.MaxInt32} {
+		key := tablecodec.EncodeRowKeyWithHandle(1, kv.IntHandle(handleId))
+		c.Assert(nextKey(key), DeepEquals, []byte(tablecodec.EncodeRowKeyWithHandle(1, kv.IntHandle(handleId+1))))
+	}
+
+	testDatums := [][]types.Datum{
+		{types.NewIntDatum(1), types.NewIntDatum(2)},
+		{types.NewIntDatum(255), types.NewIntDatum(256)},
+		{types.NewIntDatum(math.MaxInt32), types.NewIntDatum(math.MaxInt32 + 1)},
+		{types.NewStringDatum("test"), types.NewStringDatum("test\000")},
+		{types.NewStringDatum("test\255"), types.NewStringDatum("test\255\000")},
+	}
+
+	stmtCtx := new(stmtctx.StatementContext)
+	for _, datums := range testDatums {
+		keyBytes, err := codec.EncodeKey(stmtCtx, nil, types.NewIntDatum(123), datums[0])
+		c.Assert(err, IsNil)
+		h, err := kv.NewCommonHandle(keyBytes)
+		c.Assert(err, IsNil)
+		key := tablecodec.EncodeRowKeyWithHandle(1, h)
+		nextKeyBytes, err := codec.EncodeKey(stmtCtx, nil, types.NewIntDatum(123), datums[1])
+		c.Assert(err, IsNil)
+		nextHdl, err := kv.NewCommonHandle(nextKeyBytes)
+		c.Assert(err, IsNil)
+		expectNextKey := []byte(tablecodec.EncodeRowKeyWithHandle(1, nextHdl))
+		c.Assert(nextKey(key), DeepEquals, expectNextKey)
+	}
+
+	// dIAAAAAAAAD/PV9pgAAAAAD/AAABA4AAAAD/AAAAAQOAAAD/AAAAAAEAAAD8
+	// a index key with: table: 61, index: 1, int64: 1, int64: 1
+	a := []byte{116, 128, 0, 0, 0, 0, 0, 0, 255, 61, 95, 105, 128, 0, 0, 0, 0, 255, 0, 0, 1, 3, 128, 0, 0, 0, 255, 0, 0, 0, 1, 3, 128, 0, 0, 255, 0, 0, 0, 0, 1, 0, 0, 0, 252}
+	c.Assert(nextKey(a), DeepEquals, append(a, 0))
 }
 
 // The first half of this test is same as the test in tikv:
diff --git a/tests/various_types/data/vt.precise_types-schema.sql b/tests/various_types/data/vt.precise_types-schema.sql
index 04078912e..4bf75e31a 100644
--- a/tests/various_types/data/vt.precise_types-schema.sql
+++ b/tests/various_types/data/vt.precise_types-schema.sql
@@ -2,5 +2,5 @@ create table precise_types (
     a BIGINT UNSIGNED NOT NULL,
     b BIGINT NOT NULL,
     c DECIMAL(21,1) NOT NULL,
-    d DOUBLE(21,1) NOT NULL
+    d DOUBLE NOT NULL
 );
diff --git a/tests/various_types/run.sh b/tests/various_types/run.sh
index 5873a7874..fa0e72038 100755
--- a/tests/various_types/run.sh
+++ b/tests/various_types/run.sh
@@ -108,6 +108,6 @@ for BACKEND in importer tidb local; do
   check_contains 'a: 18446744073709551614'
   check_contains 'b: -9223372036854775806'
   check_contains 'c: 99999999999999999999.0'
-  check_contains 'd: 18446744073709551616.0'
+  check_contains 'd: 1.8446744073709552e19'
 
 done