// 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, util.Errorf("%T is not a range command, but EndKey is set", args) } if !rs.ContainsKey(keys.Addr(header.Key)) { return false, emptySpan, nil } return true, header, nil } // We're dealing with a range-spanning request. keyAddr, endKeyAddr := keys.Addr(header.Key), keys.Addr(header.EndKey) if l, r := !keyAddr.Equal(header.Key), !endKeyAddr.Equal(header.EndKey); l || r { if !rs.ContainsKeyRange(keyAddr, endKeyAddr) { return false, emptySpan, util.Errorf("local key range must not span ranges") } if !l || !r { return false, emptySpan, util.Errorf("local key mixed with global key in range") } // Range-local local key range. return true, header, nil } // Below, {end,}keyAddr equals header.{End,}Key, so nothing is local. if keyAddr.Less(rs.Key) { header.Key = rs.Key.AsRawKey() // "key" can't be local keyAddr = rs.Key } if !endKeyAddr.Less(rs.EndKey) { header.EndKey = rs.EndKey.AsRawKey() // "endKey" can't be local endKeyAddr = 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 !keyAddr.Less(endKeyAddr) { 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++ nReq := roachpb.RequestUnion{} if !nReq.SetValue(&roachpb.NoopRequest{}) { panic("RequestUnion excludes NoopRequest") } ba.Requests[pos] = nReq } else { // Keep the old one. If we must adjust the header, must copy. // TODO(tschottdorf): this could wind up cloning big chunks of data. // Can optimize by creating a new Request manually, but with the old // data. if newHeader.Equal(*origRequests[pos].GetInner().Header()) { ba.Requests[pos] = origRequests[pos] } else { ba.Requests[pos] = *proto.Clone(&origRequests[pos]).(*roachpb.RequestUnion) *ba.Requests[pos].GetInner().Header() = newHeader } } if err != nil { return roachpb.BatchRequest{}, 0, err } } return ba, len(ba.Requests) - numNoop, nil }