func containsKey(desc proto.RangeDescriptor, key proto.Key) bool {
return desc.ContainsKey(keys.KeyAddress(key))
}
// truncate restricts all contained requests to the given key range.
// Even on error, the returned closure must be executed; it undoes any
// truncations performed.
// First, the boundaries of the truncation are obtained: This is the
// intersection between [from,to) and the descriptor's range.
// Secondly, all requests contained in the batch are "truncated" to
// the resulting range, 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 along
// with a closure which must be executed to undo the truncation, even
// in case of an error.
// TODO(tschottdorf): Consider returning a new BatchRequest, which has more
// overhead in the common case of a batch which never needs truncation but is
// less magical.
func truncate(br *proto.BatchRequest, desc *proto.RangeDescriptor, from, to proto.Key) (func(), int, error) {
if !desc.ContainsKey(from) {
from = desc.StartKey
}
if !desc.ContainsKeyRange(desc.StartKey, to) || to == nil {
to = desc.EndKey
}
truncateOne := func(args proto.Request) (bool, []func(), error) {
if _, ok := args.(*proto.NoopRequest); ok {
return true, nil, nil
}
header := args.Header()
if !proto.IsRange(args) {
if len(header.EndKey) > 0 {
return false, nil, util.Errorf("%T is not a range command, but EndKey is set", args)
}
if !desc.ContainsKey(keys.KeyAddress(header.Key)) {
return true, nil, nil
}
return false, nil, nil
}
var undo []func()
key, endKey := header.Key, header.EndKey
keyAddr, endKeyAddr := keys.KeyAddress(key), keys.KeyAddress(endKey)
if keyAddr.Less(from) {
undo = append(undo, func() { header.Key = key })
header.Key = from
keyAddr = from
}
if !endKeyAddr.Less(to) {
undo = append(undo, func() { header.EndKey = endKey })
header.EndKey = to
endKeyAddr = to
}
// Check whether the truncation has left any keys in the range. If not,
// we need to cut it out of the request.
return !keyAddr.Less(endKeyAddr), undo, nil
}
var fns []func()
gUndo := func() {
for _, f := range fns {
f()
}
}
var numNoop int
for pos, arg := range br.Requests {
omit, undo, err := truncateOne(arg.GetValue().(proto.Request))
if omit {
numNoop++
nReq := &proto.RequestUnion{}
nReq.SetValue(&proto.NoopRequest{})
oReq := br.Requests[pos]
br.Requests[pos] = *nReq
posCpy := pos // for closure
undo = append(undo, func() {
br.Requests[posCpy] = oReq
})
}
fns = append(fns, undo...)
if err != nil {
return gUndo, 0, err
}
}
return gUndo, len(br.Requests) - numNoop, nil
}
