// Range returns a key range encompassing all the keys in the Batch.
func Range(ba roachpb.BatchRequest) (roachpb.RSpan, error) {
from := roachpb.RKeyMax
to := roachpb.RKeyMin
for _, arg := range ba.Requests {
req := arg.GetInner()
if _, ok := req.(*roachpb.NoopRequest); ok {
continue
}
h := req.Header()
if !roachpb.IsRange(req) && len(h.EndKey) != 0 {
return roachpb.RSpan{}, errors.Errorf("end key specified for non-range operation: %s", req)
}
key, err := Addr(h.Key)
if err != nil {
return roachpb.RSpan{}, err
}
if key.Less(from) {
// Key is smaller than `from`.
from = key
}
if !key.Less(to) {
// Key.Next() is larger than `to`.
if bytes.Compare(key, roachpb.RKeyMax) > 0 {
return roachpb.RSpan{}, errors.Errorf("%s must be less than KeyMax", key)
}
to = key.Next()
}
if len(h.EndKey) == 0 {
continue
}
endKey, err := AddrUpperBound(h.EndKey)
if err != nil {
return roachpb.RSpan{}, err
}
if bytes.Compare(roachpb.RKeyMax, endKey) < 0 {
return roachpb.RSpan{}, errors.Errorf("%s must be less than or equal to KeyMax", endKey)
}
if to.Less(endKey) {
// EndKey is larger than `to`.
to = endKey
}
}
return roachpb.RSpan{Key: from, EndKey: to}, nil
}
func TestKVDBInternalMethods(t *testing.T) {
defer leaktest.AfterTest(t)()
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
testCases := []roachpb.Request{
&roachpb.HeartbeatTxnRequest{},
&roachpb.GCRequest{},
&roachpb.PushTxnRequest{},
&roachpb.ResolveIntentRequest{},
&roachpb.ResolveIntentRangeRequest{},
&roachpb.MergeRequest{},
&roachpb.TruncateLogRequest{},
&roachpb.RequestLeaseRequest{},
&roachpb.EndTransactionRequest{
InternalCommitTrigger: &roachpb.InternalCommitTrigger{},
},
}
// Verify internal methods experience bad request errors.
db := createTestClient(t, s)
for i, args := range testCases {
{
header := args.Header()
header.Key = roachpb.Key("a")
args.SetHeader(header)
}
if roachpb.IsRange(args) {
header := args.Header()
header.EndKey = args.Header().Key.Next()
args.SetHeader(header)
}
b := &client.Batch{}
b.AddRawRequest(args)
err := db.Run(context.TODO(), b)
if !testutils.IsError(err, "contains an internal request|contains commit trigger") {
t.Errorf("%d: unexpected error for %s: %v", i, args.Method(), err)
}
}
}
// 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
truncBA := ba
truncBA.Requests = make([]roachpb.RequestUnion, len(ba.Requests))
for pos, arg := range ba.Requests {
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)
truncBA.Requests[pos] = union
} else {
// Keep the old one. If we must adjust the header, must copy.
if inner := ba.Requests[pos].GetInner(); newHeader.Equal(inner.Header()) {
truncBA.Requests[pos] = ba.Requests[pos]
} else {
shallowCopy := inner.ShallowCopy()
shallowCopy.SetHeader(newHeader)
union := &truncBA.Requests[pos] // avoid operating on copy
union.MustSetInner(shallowCopy)
}
}
if err != nil {
return roachpb.BatchRequest{}, 0, err
}
}
return truncBA, len(ba.Requests) - numNoop, nil
}
// fillSkippedResponses after meeting the batch key max limit for range
// requests.
func fillSkippedResponses(ba roachpb.BatchRequest, br *roachpb.BatchResponse, nextKey roachpb.RKey) {
// Some requests might have NoopResponses; we must replace them with empty
// responses of the proper type.
for i, req := range ba.Requests {
if _, ok := br.Responses[i].GetInner().(*roachpb.NoopResponse); !ok {
continue
}
var reply roachpb.Response
switch t := req.GetInner().(type) {
case *roachpb.ScanRequest:
reply = &roachpb.ScanResponse{}
case *roachpb.ReverseScanRequest:
reply = &roachpb.ReverseScanResponse{}
case *roachpb.DeleteRangeRequest:
reply = &roachpb.DeleteRangeResponse{}
case *roachpb.BeginTransactionRequest, *roachpb.EndTransactionRequest:
continue
default:
panic(fmt.Sprintf("bad type %T", t))
}
union := roachpb.ResponseUnion{}
union.MustSetInner(reply)
br.Responses[i] = union
}
// Set the ResumeSpan for future batch requests.
isReverse := ba.IsReverse()
for i, resp := range br.Responses {
req := ba.Requests[i].GetInner()
if !roachpb.IsRange(req) {
continue
}
hdr := resp.GetInner().Header()
origSpan := req.Header()
if isReverse {
if hdr.ResumeSpan != nil {
// The ResumeSpan.Key might be set to the StartKey of a range;
// correctly set it to the Key of the original request span.
hdr.ResumeSpan.Key = origSpan.Key
} else if roachpb.RKey(origSpan.Key).Less(nextKey) {
// Some keys have yet to be processed.
hdr.ResumeSpan = &origSpan
if nextKey.Less(roachpb.RKey(origSpan.EndKey)) {
// The original span has been partially processed.
hdr.ResumeSpan.EndKey = nextKey.AsRawKey()
}
}
} else {
if hdr.ResumeSpan != nil {
// The ResumeSpan.EndKey might be set to the EndKey of a
// range; correctly set it to the EndKey of the original
// request span.
hdr.ResumeSpan.EndKey = origSpan.EndKey
} else if nextKey.Less(roachpb.RKey(origSpan.EndKey)) {
// Some keys have yet to be processed.
hdr.ResumeSpan = &origSpan
if roachpb.RKey(origSpan.Key).Less(nextKey) {
// The original span has been partially processed.
hdr.ResumeSpan.Key = nextKey.AsRawKey()
}
}
}
br.Responses[i].GetInner().SetHeader(hdr)
}
}