fix bug

Signed-off-by: lance6716 <lance6716@gmail.com>
pingcap · Feb 2, 2023 · 8feb119 · 8feb119
1 parent 17996b6
commit 8feb119
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/br/pkg/lightning/backend/local/local.go b/br/pkg/lightning/backend/local/local.go
@@ -2066,7 +2066,8 @@ func nextKey(key []byte) []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) {
+ // we only do this for IntHandle, which is checked by length
+ if tablecodec.IsRecordKey(key) && len(key) == tablecodec.RecordRowKeyLen {
  tableID, handle, _ := tablecodec.DecodeRecordKey(key)
  nextHandle := handle.Next()
  // int handle overflow, use the next table prefix as nextKey
@@ -2076,7 +2077,7 @@ func nextKey(key []byte) []byte {
  return tablecodec.EncodeRowKeyWithHandle(tableID, nextHandle)
  }
 
- // if key is an index, directly append a 0x00 to the key.
+ // for index key and CommonHandle, 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/br/pkg/lightning/backend/local/local_test.go b/br/pkg/lightning/backend/local/local_test.go
@@ -115,10 +115,21 @@ func TestNextKey(t *testing.T) {
  require.NoError(t, err)
  nextHdl, err := tidbkv.NewCommonHandle(nextKeyBytes)
  require.NoError(t, err)
- expectNextKey := []byte(tablecodec.EncodeRowKeyWithHandle(1, nextHdl))
- require.Equal(t, expectNextKey, nextKey(key))
+ nextValidKey := []byte(tablecodec.EncodeRowKeyWithHandle(1, nextHdl))
+ // nextKey may return a key that can't be decoded, but it must not be larger than the valid next key.
+ require.True(t, bytes.Compare(nextKey(key), nextValidKey) <= 0, "datums: %v", datums)
  }
 
+ // a special case that when len(string datum) % 8 == 7, nextKey twice should not panic.
+ keyBytes, err := codec.EncodeKey(stmtCtx, nil, types.NewStringDatum("1234567"))
+ require.NoError(t, err)
+ h, err := tidbkv.NewCommonHandle(keyBytes)
+ require.NoError(t, err)
+ key = tablecodec.EncodeRowKeyWithHandle(1, h)
+ nextOnce := nextKey(key)
+ // should not panic
+ _ = nextKey(nextOnce)
+
  // 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}

diff --git a/br/tests/lightning_local_backend/run.sh b/br/tests/lightning_local_backend/run.sh
@@ -42,7 +42,6 @@ run_sql 'DROP DATABASE IF EXISTS cpeng;'
 run_lightning --backend local --enable-checkpoint=1 --log-file "$TEST_DIR/lightning-local.log" --config "tests/$TEST_NAME/config.toml" -L debug
 grep -Eq "split regions.*retryable error" "$TEST_DIR/lightning-local.log"
 
-read -p 123
 # Check that everything is correctly imported
 run_sql 'SELECT count(*), sum(c) FROM cpeng.a'
 check_contains 'count(*): 4'