store/tikv: Extract delete range logic from gcworker (pingcap#6093)

The deleteRanges method of gcworker contains the logic that deletes a single range. It should be separated from deleteRanges method, so we can call it somewhere else, such as testing.

This part of logic deletes all key-value pairs in a given range, which may span among several regions.

store/tikv/delete_range.go

@@ -0,0 +1,122 @@
+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikv
+
+import (
+	"bytes"
+	"context"
+
+	"github.com/juju/errors"
+	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/pingcap/tidb/store/tikv/tikvrpc"
+)
+
+// DeleteRangeTask is used to delete all keys in a range. After
+// performing DeleteRange, it keeps how many ranges it affects and
+// if the task was canceled or not.
+type DeleteRangeTask struct {
+	completedRegions int
+	canceled         bool
+	store            Storage
+	ctx              context.Context
+	bo               *Backoffer
+	startKey         []byte
+	endKey           []byte
+}
+
+// NewDeleteRangeTask creates a DeleteRangeTask. Deleting will not be performed right away.
+func NewDeleteRangeTask(ctx context.Context, store Storage, bo *Backoffer, startKey []byte, endKey []byte) *DeleteRangeTask {
+	return &DeleteRangeTask{
+		completedRegions: 0,
+		canceled:         false,
+		store:            store,
+		ctx:              ctx,
+		bo:               bo,
+		startKey:         startKey,
+		endKey:           endKey,
+	}
+}
+
+// Execute performs the delete range operation.
+func (t *DeleteRangeTask) Execute() error {
+	startKey, rangeEndKey := t.startKey, t.endKey
+	for {
+		select {
+		case <-t.ctx.Done():
+			t.canceled = true
+			return nil
+		default:
+		}
+
+		loc, err := t.store.GetRegionCache().LocateKey(t.bo, startKey)
+		if err != nil {
+			return errors.Trace(err)
+		}
+
+		// Delete to the end of the region, except if it's the last region overlapping the range
+		endKey := loc.EndKey
+		// If it is the last region
+		if loc.Contains(rangeEndKey) {
+			endKey = rangeEndKey
+		}
+
+		req := &tikvrpc.Request{
+			Type: tikvrpc.CmdDeleteRange,
+			DeleteRange: &kvrpcpb.DeleteRangeRequest{
+				StartKey: startKey,
+				EndKey:   endKey,
+			},
+		}
+
+		resp, err := t.store.SendReq(t.bo, req, loc.Region, ReadTimeoutMedium)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		regionErr, err := resp.GetRegionError()
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if regionErr != nil {
+			err = t.bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
+			if err != nil {
+				return errors.Trace(err)
+			}
+			continue
+		}
+		deleteRangeResp := resp.DeleteRange
+		if deleteRangeResp == nil {
+			return errors.Trace(ErrBodyMissing)
+		}
+		if err := deleteRangeResp.GetError(); err != "" {
+			return errors.Errorf("unexpected delete range err: %v", err)
+		}
+		t.completedRegions++
+		if bytes.Equal(endKey, rangeEndKey) {
+			break
+		}
+		startKey = endKey
+	}
+
+	return nil
+}
+
+// CompletedRegions returns the number of regions that are affected by this delete range task
+func (t *DeleteRangeTask) CompletedRegions() int {
+	return t.completedRegions
+}
+
+// IsCanceled returns true if the delete range operation was canceled on the half way
+func (t *DeleteRangeTask) IsCanceled() bool {
+	return t.canceled
+}

store/tikv/delete_range_test.go

@@ -0,0 +1,137 @@
+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikv
+
+import (
+	"bytes"
+	"context"
+	"math/rand"
+	"sort"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/tidb/store/mockstore/mocktikv"
+)
+
+type testDeleteRangeSuite struct {
+	OneByOneSuite
+	cluster *mocktikv.Cluster
+	store   *tikvStore
+}
+
+var _ = Suite(&testDeleteRangeSuite{})
+
+func (s *testDeleteRangeSuite) SetUpTest(c *C) {
+	s.cluster = mocktikv.NewCluster()
+	mocktikv.BootstrapWithMultiRegions(s.cluster, []byte("a"), []byte("b"), []byte("c"))
+	client, pdClient, err := mocktikv.NewTestClient(s.cluster, nil, "")
+	c.Assert(err, IsNil)
+
+	store, err := NewTestTiKVStore(client, pdClient, nil, nil)
+	c.Check(err, IsNil)
+	s.store = store.(*tikvStore)
+}
+
+func (s *testDeleteRangeSuite) TearDownTest(c *C) {
+	err := s.store.Close()
+	c.Assert(err, IsNil)
+}
+
+func (s *testDeleteRangeSuite) checkData(c *C, expectedData map[string]string) {
+	txn, err := s.store.Begin()
+	c.Assert(err, IsNil)
+	it, err := txn.Seek([]byte("a"))
+	c.Assert(err, IsNil)
+
+	// Scan all data and save into a map
+	data := map[string]string{}
+	for it.Valid() {
+		data[string(it.Key())] = string(it.Value())
+		err = it.Next()
+		c.Assert(err, IsNil)
+	}
+	txn.Commit(context.Background())
+
+	// Print log
+	var actualKeys []string
+	var expectedKeys []string
+	for key := range data {
+		actualKeys = append(actualKeys, key)
+	}
+	for key := range expectedData {
+		expectedKeys = append(expectedKeys, key)
+	}
+	sort.Strings(actualKeys)
+	sort.Strings(expectedKeys)
+	c.Log("Actual:   ", actualKeys)
+	c.Log("Expected: ", expectedKeys)
+
+	// Assert data in the store is the same as expected
+	c.Assert(data, DeepEquals, expectedData)
+}
+
+func (s *testDeleteRangeSuite) deleteRange(c *C, startKey []byte, endKey []byte) {
+	ctx := context.Background()
+	bo := NewBackoffer(ctx, 1000)
+	task := NewDeleteRangeTask(ctx, s.store, bo, startKey, endKey)
+
+	err := task.Execute()
+	c.Assert(err, IsNil)
+}
+
+// deleteRangeFromMap deletes all keys in a given range from a map
+func deleteRangeFromMap(m map[string]string, startKey []byte, endKey []byte) {
+	for keyStr := range m {
+		key := []byte(keyStr)
+		if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 {
+			delete(m, keyStr)
+		}
+	}
+}
+
+// testDeleteRangeOnce does delete range on both the map and the storage, and assert they are equal after deleting
+func (s *testDeleteRangeSuite) mustDeleteRange(c *C, startKey []byte, endKey []byte, expected map[string]string) {
+	s.deleteRange(c, startKey, endKey)
+	deleteRangeFromMap(expected, startKey, endKey)
+	s.checkData(c, expected)
+}
+
+func (s *testDeleteRangeSuite) TestDeleteRange(c *C) {
+	// Write some key-value pairs
+	txn, err := s.store.Begin()
+	c.Assert(err, IsNil)
+
+	testData := map[string]string{}
+
+	// Generate a sequence of keys and random values
+	for _, i := range []byte("abcd") {
+		for j := byte('0'); j <= byte('9'); j++ {
+			key := []byte{byte(i), byte(j)}
+			value := []byte{byte(rand.Intn(256)), byte(rand.Intn(256))}
+			testData[string(key)] = string(value)
+			txn.Set(key, value)
+		}
+	}
+
+	err = txn.Commit(context.Background())
+	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("d0\x00"), []byte("d1\x00"), testData)
+	s.mustDeleteRange(c, []byte("c5"), []byte("d5"), testData)
+	s.mustDeleteRange(c, []byte("a"), []byte("z"), testData)
+}

store/tikv/gcworker/gc_worker.go

@@ -351,61 +351,20 @@ func (w *GCWorker) deleteRanges(ctx context.Context, safePoint uint64) error {
 	regions := 0
 	for _, r := range ranges {
 		startKey, rangeEndKey := r.Range()
-		for {
-			select {
-			case <-ctx.Done():
-				return errors.New("[gc worker] gc job canceled")
-			default:
-			}
 
-			loc, err := w.store.GetRegionCache().LocateKey(bo, startKey)
-			if err != nil {
-				return errors.Trace(err)
-			}
+		deleteRangeTask := tikv.NewDeleteRangeTask(ctx, w.store, bo, startKey, rangeEndKey)
+		err := deleteRangeTask.Execute()
 
-			endKey := loc.EndKey
-			if loc.Contains(rangeEndKey) {
-				endKey = rangeEndKey
-			}
-
-			req := &tikvrpc.Request{
-				Type: tikvrpc.CmdDeleteRange,
-				DeleteRange: &kvrpcpb.DeleteRangeRequest{
-					StartKey: startKey,
-					EndKey:   endKey,
-				},
-			}
-
-			resp, err := w.store.SendReq(bo, req, loc.Region, tikv.ReadTimeoutMedium)
-			if err != nil {
-				return errors.Trace(err)
-			}
-			regionErr, err := resp.GetRegionError()
-			if err != nil {
-				return errors.Trace(err)
-			}
-			if regionErr != nil {
-				err = bo.Backoff(tikv.BoRegionMiss, errors.New(regionErr.String()))
-				if err != nil {
-					return errors.Trace(err)
-				}
-				continue
-			}
-			deleteRangeResp := resp.DeleteRange
-			if deleteRangeResp == nil {
-				return errors.Trace(tikv.ErrBodyMissing)
-			}
-			if err := deleteRangeResp.GetError(); err != "" {
-				return errors.Errorf("unexpected delete range err: %v", err)
-			}
-			regions++
-			if bytes.Equal(endKey, rangeEndKey) {
-				break
-			}
-			startKey = endKey
+		if err != nil {
+			return errors.Trace(err)
 		}
+		if deleteRangeTask.IsCanceled() {
+			return errors.New("[gc worker] gc job canceled")
+		}
+
+		regions += deleteRangeTask.CompletedRegions()
 		se := createSession(w.store)
-		err := util.CompleteDeleteRange(se, r)
+		err = util.CompleteDeleteRange(se, r)
 		se.Close()
 		if err != nil {
 			return errors.Trace(err)