store/tikv: Add raw delete range API

store/mockstore/mocktikv/mock_tikv_test.go

@@ -14,6 +14,7 @@
 package mocktikv
 
 import (
+ "bytes"
  "strings"
  "testing"
 
@@ -33,24 +34,81 @@ type testMockTiKVSuite struct {
 
 type testMarshal struct{}
 
+// testMvccStore is used to test MvccStore implementation.
+type testMvccStore struct {
+ testMockTiKVSuite
+}
+
+// testMVCCLevelDB is used to test MVCCLevelDB implementation.
+type testMVCCLevelDB struct {
+ testMockTiKVSuite
+}
+
 var (
  _ = Suite(&testMvccStore{})
  _ = Suite(&testMVCCLevelDB{})
  _ = Suite(testMarshal{})
 )
 
-// testMvccStore is used to test MvccStore implementation.
-type testMvccStore struct {
- testMockTiKVSuite
-}
-
 func (s *testMvccStore) SetUpTest(c *C) {
  s.store = NewMvccStore()
 }
 
-// testMVCCLevelDB is used to test MVCCLevelDB implementation.
-type testMVCCLevelDB struct {
- testMockTiKVSuite
+func (s *testMvccStore) mustRawGet(c *C, key []byte, value []byte) {
+ rawkv := s.store.(RawKV)
+ actualValue := rawkv.RawGet(key)
+ c.Assert(actualValue, DeepEquals, value)
+}
+
+func (s *testMvccStore) rawPut(c *C, key []byte, value []byte) {
+ rawkv := s.store.(RawKV)
+ rawkv.RawPut(key, value)
+}
+
+func (s *testMvccStore) checkRawData(c *C, expected map[string]string) {
+ rawkv := s.store.(RawKV)
+ pairs := rawkv.RawScan([]byte(""), []byte("zzzzzzzz"), len(expected)+1)
+ c.Assert(len(pairs), Equals, len(expected))
+
+ actual := map[string]string{}
+ for _, pair := range pairs {
+ actual[string(pair.Key)] = string(pair.Value)
+ }
+
+ c.Assert(actual, DeepEquals, expected)
+}
+
+func (s *testMvccStore) mustRawDeleteRange(c *C, startKey, endKey []byte, expectedData map[string]string) {
+ rawkv := s.store.(RawKV)
+ rawkv.RawDeleteRange(startKey, endKey)
+
+ for strKey := range expectedData {
+ key := []byte(strKey)
+ if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 {
+ delete(expectedData, strKey)
+ }
+ }
+
+ s.checkRawData(c, expectedData)
+}
+
+func (s *testMvccStore) TestRawDeleteRange(c *C) {
+ testData := map[string]string{}
+
+ for i := byte('a'); i <= byte('g'); i++ {
+ key := []byte{i}
+ value := []byte{'v', i}
+ s.rawPut(c, key, value)
+ testData[string(key)] = string(value)
+ }
+
+ s.checkRawData(c, testData)
+
+ s.mustRawDeleteRange(c, []byte("b"), []byte("b"), testData)
+ s.mustRawDeleteRange(c, []byte("d"), []byte("f"), testData)
+ s.mustRawDeleteRange(c, []byte("b"), []byte("c"), testData)
+ s.mustRawDeleteRange(c, []byte("c11"), []byte("c99"), testData)
+ s.mustRawDeleteRange(c, []byte("a"), []byte("z"), testData)
 }
 
 func (s *testMockTiKVSuite) SetUpTest(c *C) {

store/mockstore/mocktikv/mvcc.go

@@ -435,6 +435,7 @@ type RawKV interface {
  RawScan(startKey, endKey []byte, limit int) []Pair
  RawPut(key, value []byte)
  RawDelete(key []byte)
+ RawDeleteRange(startKey, endKey []byte)
 }
 
 // MVCCDebugger is for debugging.
@@ -780,6 +781,23 @@ func (s *MvccStore) RawDelete(key []byte) {
  s.rawkv.Delete(newRawEntry(key))
 }
 
+// RawDeleteRange deletes all key-value pairs in a given range
+func (s *MvccStore) RawDeleteRange(startKey, endKey []byte) {
+ s.Lock()
+ defer s.Unlock()
+
+ var entriesToDelete []*rawEntry
+ iterator := func(item btree.Item) bool {
+ entriesToDelete = append(entriesToDelete, item.(*rawEntry))
+ return true
+ }
+ s.rawkv.AscendRange(newRawEntry(startKey), newRawEntry(endKey), iterator)
+
+ for _, entry := range entriesToDelete {
+ s.rawkv.Delete(entry)
+ }
+}
+
 // RawScan reads up to a limited number of rawkv Pairs.
 func (s *MvccStore) RawScan(startKey, endKey []byte, limit int) []Pair {
  s.RLock()

store/mockstore/mocktikv/rpc.go

@@ -387,6 +387,17 @@ func (h *rpcHandler) handleKvRawDelete(req *kvrpcpb.RawDeleteRequest) *kvrpcpb.R
  return &kvrpcpb.RawDeleteResponse{}
 }
 
+func (h *rpcHandler) handleKvRawDeleteRange(req *kvrpcpb.RawDeleteRangeRequest) *kvrpcpb.RawDeleteRangeResponse {
+ rawKV, ok := h.mvccStore.(RawKV)
+ if !ok {
+ return &kvrpcpb.RawDeleteRangeResponse{
+ Error: "not implemented",
+ }
+ }
+ rawKV.RawDeleteRange(req.GetStartKey(), req.GetEndKey())
+ return &kvrpcpb.RawDeleteRangeResponse{}
+}
+
 func (h *rpcHandler) handleKvRawScan(req *kvrpcpb.RawScanRequest) *kvrpcpb.RawScanResponse {
  rawKV, ok := h.mvccStore.(RawKV)
  if !ok {
@@ -598,6 +609,13 @@ func (c *RPCClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.R
  return resp, nil
  }
  resp.RawDelete = handler.handleKvRawDelete(r)
+ case tikvrpc.CmdRawDeleteRange:
+ r := req.RawDeleteRange
+ if err := handler.checkRequest(reqCtx, r.Size()); err != nil {
+ resp.RawDeleteRange = &kvrpcpb.RawDeleteRangeResponse{RegionError: err}
+ return resp, nil
+ }
+ resp.RawDeleteRange = handler.handleKvRawDeleteRange(r)
  case tikvrpc.CmdRawScan:
  r := req.RawScan
  if err := handler.checkRequest(reqCtx, r.Size()); err != nil {

store/tikv/rawkv.go

@@ -14,6 +14,7 @@
 package tikv
 
 import (
+ "bytes"
  "time"
 
  "github.com/juju/errors"
@@ -155,6 +156,39 @@ func (c *RawKVClient) Delete(key []byte) error {
  return nil
 }
 
+// DeleteRange deletes all key-value pairs in a range from TiKV
+func (c *RawKVClient) DeleteRange(startKey []byte, endKey []byte) error {
+ start := time.Now()
+ var err error
+ defer func() {
+ var label = "delete_range"
+ if err != nil {
+ label += "_error"
+ }
+ metrics.TiKVRawkvCmdHistogram.WithLabelValues(label).Observe(time.Since(start).Seconds())
+ }()
+
+ // Process each affected region respectively
+ for !bytes.Equal(startKey, endKey) {
+ var resp *tikvrpc.Response
+ var actualEndKey []byte
+ resp, actualEndKey, err = c.sendDeleteRangeReq(startKey, endKey)
+ if err != nil {
+ return errors.Trace(err)
+ }
+ cmdResp := resp.RawDeleteRange
+ if cmdResp == nil {
+ return errors.Trace(ErrBodyMissing)
+ }
+ if cmdResp.GetError() != "" {
+ return errors.New(cmdResp.GetError())
+ }
+ startKey = actualEndKey
+ }
+
+ return nil
+}
+
 // Scan queries continuous kv pairs, starts from startKey, up to limit pairs.
 // If you want to exclude the startKey, append a '\0' to the key: `Scan(append(startKey, '\0'), limit)`.
 func (c *RawKVClient) Scan(startKey []byte, limit int) (keys [][]byte, values [][]byte, err error) {
@@ -219,3 +253,48 @@ func (c *RawKVClient) sendReq(key []byte, req *tikvrpc.Request) (*tikvrpc.Respon
  return resp, loc, nil
  }
 }
+
+// sendDeleteRangeReq sends a raw delete range request and returns the response and the actual endKey.
+// If the given range spans over more than one regions, the actual endKey is the end of the first region.
+// We can't use sendReq directly, because we need to know the end of the region before we send the request
+// TODO: Is there any better way to avoid duplicating code with func `sendReq` ?
+func (c *RawKVClient) sendDeleteRangeReq(startKey []byte, endKey []byte) (*tikvrpc.Response, []byte, error) {
+ bo := NewBackoffer(context.Background(), rawkvMaxBackoff)
+ sender := NewRegionRequestSender(c.regionCache, c.rpcClient)
+ for {
+ loc, err := c.regionCache.LocateKey(bo, startKey)
+ if err != nil {
+ return nil, nil, errors.Trace(err)
+ }
+
+ actualEndKey := endKey
+ if len(loc.EndKey) > 0 && bytes.Compare(loc.EndKey, endKey) < 0 {
+ actualEndKey = loc.EndKey
+ }
+
+ req := &tikvrpc.Request{
+ Type: tikvrpc.CmdRawDeleteRange,
+ RawDeleteRange: &kvrpcpb.RawDeleteRangeRequest{
+ StartKey: startKey,
+ EndKey: actualEndKey,
+ },
+ }
+
+ resp, err := sender.SendReq(bo, req, loc.Region, readTimeoutShort)
+ if err != nil {
+ return nil, nil, errors.Trace(err)
+ }
+ regionErr, err := resp.GetRegionError()
+ if err != nil {
+ return nil, nil, errors.Trace(err)
+ }
+ if regionErr != nil {
+ err := bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
+ if err != nil {
+ return nil, nil, errors.Trace(err)
+ }
+ continue
+ }
+ return resp, actualEndKey, nil
+ }
+}

store/tikv/rawkv_test.go

@@ -14,6 +14,8 @@
 package tikv
 
 import (
+ "bytes"
+
  . "github.com/pingcap/check"
  "github.com/pingcap/tidb/store/mockstore/mocktikv"
  "golang.org/x/net/context"
@@ -78,6 +80,41 @@ func (s *testRawKVSuite) mustScan(c *C, startKey string, limit int, expect ...st
  }
 }
 
+func (s *testRawKVSuite) mustDeleteRange(c *C, startKey, endKey []byte, expected map[string]string) {
+ err := s.client.DeleteRange(startKey, endKey)
+ c.Assert(err, IsNil)
+
+ for keyStr := range expected {
+ key := []byte(keyStr)
+ if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 {
+ delete(expected, keyStr)
+ }
+ }
+
+ s.checkData(c, expected)
+}
+
+func (s *testRawKVSuite) checkData(c *C, expected map[string]string) {
+ keys, values, err := s.client.Scan([]byte(""), len(expected)+1)
+ c.Assert(err, IsNil)
+
+ c.Assert(len(expected), Equals, len(keys))
+ for i, key := range keys {
+ c.Assert(expected[string(key)], Equals, string(values[i]))
+ }
+}
+
+func (s *testRawKVSuite) split(c *C, regionKey, splitKey string) error {
+ loc, err := s.client.regionCache.LocateKey(s.bo, []byte(regionKey))
+ if err != nil {
+ return err
+ }
+
+ newRegionID, peerID := s.cluster.AllocID(), s.cluster.AllocID()
+ s.cluster.SplitRaw(loc.Region.id, newRegionID, []byte(splitKey), []uint64{peerID}, peerID)
+ return nil
+}
+
 func (s *testRawKVSuite) TestSimple(c *C) {
  s.mustNotExist(c, []byte("key"))
  s.mustPut(c, []byte("key"), []byte("value"))
@@ -89,13 +126,11 @@ func (s *testRawKVSuite) TestSimple(c *C) {
 }
 
 func (s *testRawKVSuite) TestSplit(c *C) {
- loc, err := s.client.regionCache.LocateKey(s.bo, []byte("k"))
- c.Assert(err, IsNil)
  s.mustPut(c, []byte("k1"), []byte("v1"))
  s.mustPut(c, []byte("k3"), []byte("v3"))
 
- newRegionID, peerID := s.cluster.AllocID(), s.cluster.AllocID()
- s.cluster.SplitRaw(loc.Region.id, newRegionID, []byte("k2"), []uint64{peerID}, peerID)
+ err := s.split(c, "k", "k2")
+ c.Assert(err, IsNil)
 
  s.mustGet(c, []byte("k1"), []byte("v1"))
  s.mustGet(c, []byte("k3"), []byte("v3"))
@@ -115,16 +150,41 @@ func (s *testRawKVSuite) TestScan(c *C) {
  s.mustScan(c, "k2", 3, "k3", "v3", "k5", "v5", "k7", "v7")
  }
 
- split := func(regionKey, splitKey string) {
- loc, err := s.client.regionCache.LocateKey(s.bo, []byte(regionKey))
- c.Assert(err, IsNil)
- newRegionID, peerID := s.cluster.AllocID(), s.cluster.AllocID()
- s.cluster.SplitRaw(loc.Region.id, newRegionID, []byte(splitKey), []uint64{peerID}, peerID)
- }
-
  check()
- split("k", "k2")
+
+ err := s.split(c, "k", "k2")
+ c.Assert(err, IsNil)
  check()
- split("k2", "k5")
+
+ err = s.split(c, "k2", "k5")
+ c.Assert(err, IsNil)
  check()
 }
+
+func (s *testRawKVSuite) TestDeleteRange(c *C) {
+ // Init data
+ testData := map[string]string{}
+ for _, i := range []byte("abcd") {
+ for j := byte('0'); j <= byte('9'); j++ {
+ key := []byte{i, j}
+ value := []byte{'v', i, j}
+ s.mustPut(c, key, value)
+
+ testData[string(key)] = string(value)
+ }
+ }
+
+ err := s.split(c, "b", "b")
+ c.Assert(err, IsNil)
+ err = s.split(c, "c", "c")
+ c.Assert(err, IsNil)
+ err = s.split(c, "d", "d")
+ c.Assert(err, IsNil)
+
+ s.checkData(c, testData)
+ s.mustDeleteRange(c, []byte("b"), []byte("c0"), testData)
+ s.mustDeleteRange(c, []byte("c11"), []byte("c12"), testData)
+ s.mustDeleteRange(c, []byte("d0"), []byte("d0"), testData)
+ s.mustDeleteRange(c, []byte("c5"), []byte("d5"), testData)
+ s.mustDeleteRange(c, []byte("a"), []byte("z"), testData)
+}