/
pd_codec_v2.go
110 lines (92 loc) · 3.61 KB
/
pd_codec_v2.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package locate
import (
"context"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pingcap/kvproto/pkg/pdpb"
"github.com/tikv/client-go/v2/internal/client"
"github.com/tikv/client-go/v2/util/codec"
pd "github.com/tikv/pd/client"
)
// CodecPDClientV2 wraps a PD Client to decode the region meta in API v2 manner.
type CodecPDClientV2 struct {
*CodecPDClient
mode client.Mode
}
// NewCodecPDClientV2 create a CodecPDClientV2.
func NewCodecPDClientV2(client pd.Client, mode client.Mode) *CodecPDClientV2 {
codecClient := NewCodeCPDClient(client)
return &CodecPDClientV2{codecClient, mode}
}
// GetRegion encodes the key before send requests to pd-server and decodes the
// returned StartKey && EndKey from pd-server.
func (c *CodecPDClientV2) GetRegion(ctx context.Context, key []byte, opts ...pd.GetRegionOption) (*pd.Region, error) {
queryKey := client.EncodeV2Key(c.mode, key)
region, err := c.CodecPDClient.GetRegion(ctx, queryKey, opts...)
return c.processRegionResult(region, err)
}
// GetPrevRegion encodes the key before send requests to pd-server and decodes the
// returned StartKey && EndKey from pd-server.
func (c *CodecPDClientV2) GetPrevRegion(ctx context.Context, key []byte, opts ...pd.GetRegionOption) (*pd.Region, error) {
queryKey := client.EncodeV2Key(c.mode, key)
region, err := c.CodecPDClient.GetPrevRegion(ctx, queryKey, opts...)
return c.processRegionResult(region, err)
}
// GetRegionByID encodes the key before send requests to pd-server and decodes the
// returned StartKey && EndKey from pd-server.
func (c *CodecPDClientV2) GetRegionByID(ctx context.Context, regionID uint64, opts ...pd.GetRegionOption) (*pd.Region, error) {
region, err := c.CodecPDClient.GetRegionByID(ctx, regionID, opts...)
return c.processRegionResult(region, err)
}
// ScanRegions encodes the key before send requests to pd-server and decodes the
// returned StartKey && EndKey from pd-server.
func (c *CodecPDClientV2) ScanRegions(ctx context.Context, startKey []byte, endKey []byte, limit int) ([]*pd.Region, error) {
start, end := client.EncodeV2Range(c.mode, startKey, endKey)
regions, err := c.CodecPDClient.ScanRegions(ctx, start, end, limit)
if err != nil {
return nil, err
}
for i := range regions {
region, _ := c.processRegionResult(regions[i], nil)
regions[i] = region
}
return regions, nil
}
// SplitRegions split regions by given split keys
func (c *CodecPDClientV2) SplitRegions(ctx context.Context, splitKeys [][]byte, opts ...pd.RegionsOption) (*pdpb.SplitRegionsResponse, error) {
var keys [][]byte
for i := range splitKeys {
withPrefix := client.EncodeV2Key(c.mode, splitKeys[i])
keys = append(keys, codec.EncodeBytes(nil, withPrefix))
}
return c.CodecPDClient.SplitRegions(ctx, keys, opts...)
}
func (c *CodecPDClientV2) processRegionResult(region *pd.Region, err error) (*pd.Region, error) {
if err != nil {
return nil, err
}
if region != nil {
// TODO(@iosmanthus): enable buckets support.
region.Buckets = nil
region.Meta.StartKey, region.Meta.EndKey =
client.MapV2RangeToV1(c.mode, region.Meta.StartKey, region.Meta.EndKey)
}
return region, nil
}
func (c *CodecPDClientV2) decodeRegionWithShallowCopy(region *metapb.Region) (*metapb.Region, error) {
var err error
newRegion := *region
if len(region.StartKey) > 0 {
_, newRegion.StartKey, err = codec.DecodeBytes(region.StartKey, nil)
}
if err != nil {
return nil, err
}
if len(region.EndKey) > 0 {
_, newRegion.EndKey, err = codec.DecodeBytes(region.EndKey, nil)
}
if err != nil {
return nil, err
}
newRegion.StartKey, newRegion.EndKey = client.MapV2RangeToV1(c.mode, newRegion.StartKey, newRegion.EndKey)
return &newRegion, nil
}