func newReplicaDataIterator(d *roachpb.RangeDescriptor, e engine.Engine) *replicaDataIterator {
// The first range in the keyspace starts at KeyMin, which includes the node-local
// space. We need the original StartKey to find the range metadata, but the
// actual data starts at LocalMax.
dataStartKey := d.StartKey.AsRawKey()
if d.StartKey.Equal(roachpb.RKeyMin) {
dataStartKey = keys.LocalMax
}
ri := &replicaDataIterator{
ranges: []keyRange{
{
start: engine.MVCCEncodeKey(keys.MakeRangeIDPrefix(d.RangeID)),
end: engine.MVCCEncodeKey(keys.MakeRangeIDPrefix(d.RangeID + 1)),
},
{
start: engine.MVCCEncodeKey(keys.MakeRangeKeyPrefix(d.StartKey)),
end: engine.MVCCEncodeKey(keys.MakeRangeKeyPrefix(d.EndKey)),
},
{
start: engine.MVCCEncodeKey(dataStartKey),
end: engine.MVCCEncodeKey(d.EndKey.AsRawKey()),
},
},
iter: e.NewIterator(),
}
ri.iter.Seek(ri.ranges[ri.curIndex].start)
ri.advance()
return ri
}
func makeReplicaKeyRanges(d *roachpb.RangeDescriptor) []keyRange {
// The first range in the keyspace starts at KeyMin, which includes the
// node-local space. We need the original StartKey to find the range
// metadata, but the actual data starts at LocalMax.
dataStartKey := d.StartKey.AsRawKey()
if d.StartKey.Equal(roachpb.RKeyMin) {
dataStartKey = keys.LocalMax
}
return []keyRange{
{
start: engine.MakeMVCCMetadataKey(keys.MakeRangeIDPrefix(d.RangeID)),
end: engine.MakeMVCCMetadataKey(keys.MakeRangeIDPrefix(d.RangeID + 1)),
},
{
start: engine.MakeMVCCMetadataKey(keys.MakeRangeKeyPrefix(d.StartKey)),
end: engine.MakeMVCCMetadataKey(keys.MakeRangeKeyPrefix(d.EndKey)),
},
{
start: engine.MakeMVCCMetadataKey(dataStartKey),
end: engine.MakeMVCCMetadataKey(d.EndKey.AsRawKey()),
},
}
}
func TestTruncate(t *testing.T) {
defer leaktest.AfterTest(t)()
loc := func(s string) string {
return string(keys.RangeDescriptorKey(roachpb.RKey(s)))
}
locPrefix := func(s string) string {
return string(keys.MakeRangeKeyPrefix(roachpb.RKey(s)))
}
testCases := []struct {
keys [][2]string
expKeys [][2]string
from, to string
desc [2]string // optional, defaults to {from,to}
err string
}{
{
// Keys inside of active range.
keys: [][2]string{{"a", "q"}, {"c"}, {"b, e"}, {"q"}},
expKeys: [][2]string{{"a", "q"}, {"c"}, {"b, e"}, {"q"}},
from: "a", to: "q\x00",
},
{
// Keys outside of active range.
keys: [][2]string{{"a"}, {"a", "b"}, {"q"}, {"q", "z"}},
expKeys: [][2]string{{}, {}, {}, {}},
from: "b", to: "q",
},
{
// Range-local keys inside of active range.
keys: [][2]string{{loc("b")}, {loc("c")}},
expKeys: [][2]string{{loc("b")}, {loc("c")}},
from: "b", to: "e",
},
{
// Range-local key outside of active range.
keys: [][2]string{{loc("a")}},
expKeys: [][2]string{{}},
from: "b", to: "e",
},
{
// Range-local range contained in active range.
keys: [][2]string{{loc("b"), loc("e") + "\x00"}},
expKeys: [][2]string{{loc("b"), loc("e") + "\x00"}},
from: "b", to: "e\x00",
},
{
// Range-local range not contained in active range.
keys: [][2]string{{loc("a"), loc("b")}},
expKeys: [][2]string{{}},
from: "c", to: "e",
},
{
// Range-local range not contained in active range.
keys: [][2]string{{loc("a"), locPrefix("b")}, {loc("e"), loc("f")}},
expKeys: [][2]string{{}, {}},
from: "b", to: "e",
},
{
// Range-local range partially contained in active range.
keys: [][2]string{{loc("a"), loc("b")}},
expKeys: [][2]string{{loc("a"), locPrefix("b")}},
from: "a", to: "b",
},
{
// Range-local range partially contained in active range.
keys: [][2]string{{loc("a"), loc("b")}},
expKeys: [][2]string{{locPrefix("b"), loc("b")}},
from: "b", to: "e",
},
{
// Range-local range contained in active range.
keys: [][2]string{{locPrefix("b"), loc("b")}},
expKeys: [][2]string{{locPrefix("b"), loc("b")}},
from: "b", to: "c",
},
{
// Mixed range-local vs global key range.
keys: [][2]string{{loc("c"), "d\x00"}},
from: "b", to: "e",
err: "local key mixed with global key",
},
{
// Key range touching and intersecting active range.
keys: [][2]string{{"a", "b"}, {"a", "c"}, {"p", "q"}, {"p", "r"}, {"a", "z"}},
expKeys: [][2]string{{}, {"b", "c"}, {"p", "q"}, {"p", "q"}, {"b", "q"}},
from: "b", to: "q",
},
// Active key range is intersection of descriptor and [from,to).
{
keys: [][2]string{{"c", "q"}},
expKeys: [][2]string{{"d", "p"}},
from: "a", to: "z",
desc: [2]string{"d", "p"},
},
{
keys: [][2]string{{"c", "q"}},
expKeys: [][2]string{{"d", "p"}},
from: "d", to: "p",
desc: [2]string{"a", "z"},
},
}
for i, test := range testCases {
goldenOriginal := roachpb.BatchRequest{}
for _, ks := range test.keys {
if len(ks[1]) > 0 {
goldenOriginal.Add(&roachpb.ResolveIntentRangeRequest{
Span: roachpb.Span{Key: roachpb.Key(ks[0]), EndKey: roachpb.Key(ks[1])},
IntentTxn: roachpb.TxnMeta{ID: uuid.NewV4()},
})
} else {
goldenOriginal.Add(&roachpb.GetRequest{
Span: roachpb.Span{Key: roachpb.Key(ks[0])},
})
}
}
original := roachpb.BatchRequest{Requests: make([]roachpb.RequestUnion, len(goldenOriginal.Requests))}
for i, request := range goldenOriginal.Requests {
original.Requests[i].SetValue(request.GetInner().ShallowCopy())
}
desc := &roachpb.RangeDescriptor{
StartKey: roachpb.RKey(test.desc[0]), EndKey: roachpb.RKey(test.desc[1]),
}
if len(desc.StartKey) == 0 {
desc.StartKey = roachpb.RKey(test.from)
}
if len(desc.EndKey) == 0 {
desc.EndKey = roachpb.RKey(test.to)
}
rs := roachpb.RSpan{Key: roachpb.RKey(test.from), EndKey: roachpb.RKey(test.to)}
rs, err := rs.Intersect(desc)
if err != nil {
t.Errorf("%d: intersection failure: %v", i, err)
continue
}
ba, num, err := truncate(original, rs)
if err != nil || test.err != "" {
if test.err == "" || !testutils.IsError(err, test.err) {
t.Errorf("%d: %v (expected: %s)", i, err, test.err)
}
continue
}
var reqs int
for j, arg := range ba.Requests {
req := arg.GetInner()
if h := req.Header(); !bytes.Equal(h.Key, roachpb.Key(test.expKeys[j][0])) || !bytes.Equal(h.EndKey, roachpb.Key(test.expKeys[j][1])) {
t.Errorf("%d.%d: range mismatch: actual [%q,%q), wanted [%q,%q)", i, j,
h.Key, h.EndKey, test.expKeys[j][0], test.expKeys[j][1])
} else if _, ok := req.(*roachpb.NoopRequest); ok != (len(h.Key) == 0) {
t.Errorf("%d.%d: expected NoopRequest, got %T", i, j, req)
} else if len(h.Key) != 0 {
reqs++
}
}
if reqs != num {
t.Errorf("%d: counted %d requests, but truncation indicated %d", i, reqs, num)
}
if !reflect.DeepEqual(original, goldenOriginal) {
t.Errorf("%d: truncation mutated original:\nexpected: %s\nactual: %s",
i, goldenOriginal, original)
}
}
}
// truncate restricts all contained requests to the given key range
// and returns a new BatchRequest.
// All requests contained in that batch are "truncated" to the given
// span, inserting NoopRequest appropriately to replace requests which
// are left without a key range to operate on. The number of non-noop
// requests after truncation is returned.
func truncate(ba roachpb.BatchRequest, rs roachpb.RSpan) (roachpb.BatchRequest, int, error) {
truncateOne := func(args roachpb.Request) (bool, roachpb.Span, error) {
if _, ok := args.(*roachpb.NoopRequest); ok {
return true, emptySpan, nil
}
header := args.Header()
if !roachpb.IsRange(args) {
// This is a point request.
if len(header.EndKey) > 0 {
return false, emptySpan, errors.Errorf("%T is not a range command, but EndKey is set", args)
}
keyAddr, err := keys.Addr(header.Key)
if err != nil {
return false, emptySpan, err
}
if !rs.ContainsKey(keyAddr) {
return false, emptySpan, nil
}
return true, header, nil
}
// We're dealing with a range-spanning request.
local := false
keyAddr, err := keys.Addr(header.Key)
if err != nil {
return false, emptySpan, err
}
endKeyAddr, err := keys.Addr(header.EndKey)
if err != nil {
return false, emptySpan, err
}
if l, r := !keyAddr.Equal(header.Key), !endKeyAddr.Equal(header.EndKey); l || r {
if !l || !r {
return false, emptySpan, errors.Errorf("local key mixed with global key in range")
}
local = true
}
if keyAddr.Less(rs.Key) {
// rs.Key can't be local because it contains range split points, which
// are never local.
if !local {
header.Key = rs.Key.AsRawKey()
} else {
// The local start key should be truncated to the boundary of local keys which
// address to rs.Key.
header.Key = keys.MakeRangeKeyPrefix(rs.Key)
}
}
if !endKeyAddr.Less(rs.EndKey) {
// rs.EndKey can't be local because it contains range split points, which
// are never local.
if !local {
header.EndKey = rs.EndKey.AsRawKey()
} else {
// The local end key should be truncated to the boundary of local keys which
// address to rs.EndKey.
header.EndKey = keys.MakeRangeKeyPrefix(rs.EndKey)
}
}
// Check whether the truncation has left any keys in the range. If not,
// we need to cut it out of the request.
if header.Key.Compare(header.EndKey) >= 0 {
return false, emptySpan, nil
}
return true, header, nil
}
var numNoop int
origRequests := ba.Requests
ba.Requests = make([]roachpb.RequestUnion, len(ba.Requests))
for pos, arg := range origRequests {
hasRequest, newHeader, err := truncateOne(arg.GetInner())
if !hasRequest {
// We omit this one, i.e. replace it with a Noop.
numNoop++
union := roachpb.RequestUnion{}
union.MustSetInner(&noopRequest)
ba.Requests[pos] = union
} else {
// Keep the old one. If we must adjust the header, must copy.
if inner := origRequests[pos].GetInner(); newHeader.Equal(inner.Header()) {
ba.Requests[pos] = origRequests[pos]
} else {
shallowCopy := inner.ShallowCopy()
shallowCopy.SetHeader(newHeader)
union := &ba.Requests[pos] // avoid operating on copy
union.MustSetInner(shallowCopy)
}
}
if err != nil {
return roachpb.BatchRequest{}, 0, err
}
}
return ba, len(ba.Requests) - numNoop, nil
}
// TestTruncateWithLocalSpanAndDescriptor verifies that a batch request with local keys
// is truncated with a range span and the range of a descriptor found in cache.
func TestTruncateWithLocalSpanAndDescriptor(t *testing.T) {
defer leaktest.AfterTest(t)()
g, s := makeTestGossip(t)
defer s()
if err := g.SetNodeDescriptor(&roachpb.NodeDescriptor{NodeID: 1}); err != nil {
t.Fatal(err)
}
nd := &roachpb.NodeDescriptor{
NodeID: roachpb.NodeID(1),
Address: util.MakeUnresolvedAddr(testAddress.Network(), testAddress.String()),
}
if err := g.AddInfoProto(gossip.MakeNodeIDKey(roachpb.NodeID(1)), nd, time.Hour); err != nil {
t.Fatal(err)
}
// Fill mockRangeDescriptorDB with two descriptors.
var descriptor1 = roachpb.RangeDescriptor{
RangeID: 1,
StartKey: roachpb.RKeyMin,
EndKey: roachpb.RKey("b"),
Replicas: []roachpb.ReplicaDescriptor{
{
NodeID: 1,
StoreID: 1,
},
},
}
var descriptor2 = roachpb.RangeDescriptor{
RangeID: 2,
StartKey: roachpb.RKey("b"),
EndKey: roachpb.RKey("c"),
Replicas: []roachpb.ReplicaDescriptor{
{
NodeID: 1,
StoreID: 1,
},
},
}
var descriptor3 = roachpb.RangeDescriptor{
RangeID: 3,
StartKey: roachpb.RKey("c"),
EndKey: roachpb.RKeyMax,
Replicas: []roachpb.ReplicaDescriptor{
{
NodeID: 1,
StoreID: 1,
},
},
}
descDB := mockRangeDescriptorDB(func(key roachpb.RKey, _, _ bool) ([]roachpb.RangeDescriptor, *roachpb.Error) {
switch {
case !key.Less(roachpb.RKey("c")):
return []roachpb.RangeDescriptor{descriptor3}, nil
case !key.Less(roachpb.RKey("b")):
return []roachpb.RangeDescriptor{descriptor2}, nil
default:
return []roachpb.RangeDescriptor{descriptor1}, nil
}
})
// Define our rpcSend stub which checks the span of the batch
// requests.
requests := 0
sendStub := func(_ SendOptions, _ ReplicaSlice, ba roachpb.BatchRequest, _ *rpc.Context) (*roachpb.BatchResponse, error) {
h := ba.Requests[0].GetInner().Header()
switch requests {
case 0:
wantStart := keys.RangeDescriptorKey(roachpb.RKey("a"))
wantEnd := keys.MakeRangeKeyPrefix(roachpb.RKey("b"))
if !(h.Key.Equal(wantStart) && h.EndKey.Equal(wantEnd)) {
t.Errorf("Unexpected span [%s,%s), want [%s,%s)", h.Key, h.EndKey, wantStart, wantEnd)
}
case 1:
wantStart := keys.MakeRangeKeyPrefix(roachpb.RKey("b"))
wantEnd := keys.MakeRangeKeyPrefix(roachpb.RKey("c"))
if !(h.Key.Equal(wantStart) && h.EndKey.Equal(wantEnd)) {
t.Errorf("Unexpected span [%s,%s), want [%s,%s)", h.Key, h.EndKey, wantStart, wantEnd)
}
case 2:
wantStart := keys.MakeRangeKeyPrefix(roachpb.RKey("c"))
wantEnd := keys.RangeDescriptorKey(roachpb.RKey("c"))
if !(h.Key.Equal(wantStart) && h.EndKey.Equal(wantEnd)) {
t.Errorf("Unexpected span [%s,%s), want [%s,%s)", h.Key, h.EndKey, wantStart, wantEnd)
}
}
requests++
batchReply := &roachpb.BatchResponse{}
reply := &roachpb.ScanResponse{}
batchReply.Add(reply)
return batchReply, nil
}
ctx := &DistSenderContext{
RPCSend: sendStub,
RangeDescriptorDB: descDB,
}
ds := NewDistSender(ctx, g)
// Send a batch request contains two scans. In the first
// attempt, the range of the descriptor found in the cache is
// ["", "b"). The request is truncated to contain only the scan
// on local keys that address up to "b".
//
// In the second attempt, The range of the descriptor found in
// the cache is ["b", "d"), The request is truncated to contain
// only the scan on local keys that address from "b" to "d".
ba := roachpb.BatchRequest{}
ba.Txn = &roachpb.Transaction{Name: "test"}
ba.Add(roachpb.NewScan(keys.RangeDescriptorKey(roachpb.RKey("a")), keys.RangeDescriptorKey(roachpb.RKey("c")), 0))
if _, pErr := ds.Send(context.Background(), ba); pErr != nil {
t.Fatal(pErr)
}
if want := 3; requests != want {
t.Errorf("expected request to be split into %d parts, found %d", want, requests)
}
}