move unionstore package to internal (#179)

Signed-off-by: disksing <i@disksing.com>

integration_tests/scan_test.go

@@ -42,7 +42,6 @@ import (
 	"github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/logutil"
 	"github.com/tikv/client-go/v2/tikv"
-	"github.com/tikv/client-go/v2/unionstore"
 	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 )
@@ -104,7 +103,7 @@ func (s *testScanSuite) makeValue(i int) []byte {
 }
 
 func (s *testScanSuite) TestScan() {
-	check := func(scan unionstore.Iterator, rowNum int, keyOnly bool) {
+	check := func(scan tikv.Iterator, rowNum int, keyOnly bool) {
 		for i := 0; i < rowNum; i++ {
 			k := scan.Key()
 			expectedKey := s.makeKey(i)

tikv/2pc.go

@@ -52,13 +52,13 @@ import (
 	"github.com/tikv/client-go/v2/config"
 	tikverr "github.com/tikv/client-go/v2/error"
 	"github.com/tikv/client-go/v2/internal/locate"
+	"github.com/tikv/client-go/v2/internal/unionstore"
 	"github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/logutil"
 	"github.com/tikv/client-go/v2/metrics"
 	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/retry"
 	"github.com/tikv/client-go/v2/tikvrpc"
-	"github.com/tikv/client-go/v2/unionstore"
 	"github.com/tikv/client-go/v2/util"
 	zap "go.uber.org/zap"
 )
@@ -136,11 +136,11 @@ type twoPhaseCommitter struct {
 }
 
 type memBufferMutations struct {
-	storage *unionstore.MemDB
+	storage *MemDB
 	handles []unionstore.MemKeyHandle
 }
 
-func newMemBufferMutations(sizeHint int, storage *unionstore.MemDB) *memBufferMutations {
+func newMemBufferMutations(sizeHint int, storage *MemDB) *memBufferMutations {
 	return &memBufferMutations{
 		handles: make([]unionstore.MemKeyHandle, 0, sizeHint),
 		storage: storage,

tikv/batch_getter.go

@@ -37,13 +37,12 @@ import (
 
 	"github.com/pingcap/errors"
 	tikverr "github.com/tikv/client-go/v2/error"
-	"github.com/tikv/client-go/v2/unionstore"
 )
 
 // BatchBufferGetter is the interface for BatchGet.
 type BatchBufferGetter interface {
 	Len() int
-	unionstore.Getter
+	Getter
 }
 
 // BatchGetter is the interface for BatchGet.

tikv/snapshot.go

@@ -54,7 +54,6 @@ import (
 	"github.com/tikv/client-go/v2/metrics"
 	"github.com/tikv/client-go/v2/retry"
 	"github.com/tikv/client-go/v2/tikvrpc"
-	"github.com/tikv/client-go/v2/unionstore"
 	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 )
@@ -616,13 +615,13 @@ func (s *KVSnapshot) mergeExecDetail(detail *kvrpcpb.ExecDetailsV2) {
 }
 
 // Iter return a list of key-value pair after `k`.
-func (s *KVSnapshot) Iter(k []byte, upperBound []byte) (unionstore.Iterator, error) {
+func (s *KVSnapshot) Iter(k []byte, upperBound []byte) (Iterator, error) {
 	scanner, err := newScanner(s, k, upperBound, scanBatchSize, false)
 	return scanner, errors.Trace(err)
 }
 
 // IterReverse creates a reversed Iterator positioned on the first entry which key is less than k.
-func (s *KVSnapshot) IterReverse(k []byte) (unionstore.Iterator, error) {
+func (s *KVSnapshot) IterReverse(k []byte) (Iterator, error) {
 	scanner, err := newScanner(s, nil, k, scanBatchSize, true)
 	return scanner, errors.Trace(err)
 }

tikv/test_probe.go

@@ -42,9 +42,9 @@ import (
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/tikv/client-go/v2/internal/locate"
+	"github.com/tikv/client-go/v2/internal/unionstore"
 	"github.com/tikv/client-go/v2/retry"
 	"github.com/tikv/client-go/v2/tikvrpc"
-	"github.com/tikv/client-go/v2/unionstore"
 	pd "github.com/tikv/pd/client"
 )
 

tikv/txn.go

@@ -49,12 +49,12 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
 	tikverr "github.com/tikv/client-go/v2/error"
+	"github.com/tikv/client-go/v2/internal/unionstore"
 	tikv "github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/logutil"
 	"github.com/tikv/client-go/v2/metrics"
 	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/retry"
-	"github.com/tikv/client-go/v2/unionstore"
 	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 )
@@ -217,12 +217,12 @@ func (txn *KVTxn) String() string {
 // If such entry is not found, it returns an invalid Iterator with no error.
 // It yields only keys that < upperBound. If upperBound is nil, it means the upperBound is unbounded.
 // The Iterator must be Closed after use.
-func (txn *KVTxn) Iter(k []byte, upperBound []byte) (unionstore.Iterator, error) {
+func (txn *KVTxn) Iter(k []byte, upperBound []byte) (Iterator, error) {
 	return txn.us.Iter(k, upperBound)
 }
 
 // IterReverse creates a reversed Iterator positioned on the first entry which key is less than k.
-func (txn *KVTxn) IterReverse(k []byte) (unionstore.Iterator, error) {
+func (txn *KVTxn) IterReverse(k []byte) (Iterator, error) {
 	return txn.us.IterReverse(k)
 }
 
@@ -763,7 +763,7 @@ func (txn *KVTxn) GetUnionStore() *unionstore.KVUnionStore {
 }
 
 // GetMemBuffer return the MemBuffer binding to this transaction.
-func (txn *KVTxn) GetMemBuffer() *unionstore.MemDB {
+func (txn *KVTxn) GetMemBuffer() *MemDB {
 	return txn.us.GetMemBuffer()
 }
 

tikv/unionstore_export.go

@@ -0,0 +1,51 @@
+// Copyright 2021 TiKV Authors
+//
+// 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.
+
+// NOTE: The code in this file is based on code from the
+// TiDB project, licensed under the Apache License v 2.0
+//
+// https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/txn.go
+//
+
+// Copyright 2016 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 "github.com/tikv/client-go/v2/internal/unionstore"
+
+// Getter is the interface for the Get method.
+type Getter = unionstore.Getter
+
+// Iterator is the interface for a iterator on KV store.
+type Iterator = unionstore.Iterator
+
+// MemDB is rollbackable Red-Black Tree optimized for transaction states buffer use scenario.
+// You can think MemDB is a combination of two separate tree map, one for key => value and another for key => keyFlags.
+//
+// The value map is rollbackable, that means you can use the `Staging`, `Release` and `Cleanup` API to safely modify KVs.
+//
+// The flags map is not rollbackable. There are two types of flag, persistent and non-persistent.
+// When discarding a newly added KV in `Cleanup`, the non-persistent flags will be cleared.
+// If there are persistent flags associated with key, we will keep this key in node without value.
+type MemDB = unionstore.MemDB

unionstore/unionstore.go

@@ -0,0 +1,35 @@
+// Copyright 2021 TiKV Authors
+//
+// 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 unionstore
+
+import "github.com/tikv/client-go/v2/internal/unionstore"
+
+// Re-export internal/unionstore for compatibility.
+// TODO: Remove it.
+
+// Getter is the interface for the Get method.
+type Getter = unionstore.Getter
+
+// Iterator is the interface for a iterator on KV store.
+type Iterator = unionstore.Iterator
+
+// MemDB is rollbackable Red-Black Tree optimized for transaction states buffer use scenario.
+// You can think MemDB is a combination of two separate tree map, one for key => value and another for key => keyFlags.
+//
+// The value map is rollbackable, that means you can use the `Staging`, `Release` and `Cleanup` API to safely modify KVs.
+//
+// The flags map is not rollbackable. There are two types of flag, persistent and non-persistent.
+// When discarding a newly added KV in `Cleanup`, the non-persistent flags will be cleared.
+// If there are persistent flags associated with key, we will keep this key in node without value.
+type MemDB = unionstore.MemDB