// flush writes all dirty nodes and the tree to the transaction.
func (tc *treeContext) flush(b *client.Batch) {
if tc.dirty {
b.Put(keys.RangeTreeRoot, tc.tree)
}
for key, cachedNode := range tc.nodes {
if cachedNode.dirty {
if cachedNode.node == nil {
b.Del(keys.RangeTreeNodeKey(roachpb.RKey(key)))
} else {
b.Put(keys.RangeTreeNodeKey(roachpb.RKey(key)), cachedNode.node)
}
}
}
}
// compareBiogoNode compares a biogo node and a range tree node to determine if both
// contain the same values in the same order. It recursively calls itself on
// both children if they exist.
func compareBiogoNode(db *client.DB, biogoNode *llrb.Node, key *proto.Key) error {
// Retrieve the node form the range tree.
rtNode := &proto.RangeTreeNode{}
if err := db.GetProto(keys.RangeTreeNodeKey(*key), rtNode); err != nil {
return err
}
bNode := &proto.RangeTreeNode{
Key: proto.Key(biogoNode.Elem.(Key)),
ParentKey: proto.KeyMin,
Black: bool(biogoNode.Color),
}
if biogoNode.Left != nil {
leftKey := proto.Key(biogoNode.Left.Elem.(Key))
bNode.LeftKey = &leftKey
}
if biogoNode.Right != nil {
rightKey := proto.Key(biogoNode.Right.Elem.(Key))
bNode.RightKey = &rightKey
}
if err := nodesEqual(*key, *bNode, *rtNode); err != nil {
return err
}
if rtNode.LeftKey != nil {
if err := compareBiogoNode(db, biogoNode.Left, rtNode.LeftKey); err != nil {
return err
}
}
if rtNode.RightKey != nil {
if err := compareBiogoNode(db, biogoNode.Right, rtNode.RightKey); err != nil {
return err
}
}
return nil
}
// flush writes all dirty nodes and the tree to the transaction.
func (tc *treeContext) flush(b *client.Batch) error {
if tc.dirty {
b.Put(keys.RangeTreeRoot, tc.tree)
}
for _, cachedNode := range tc.nodes {
if cachedNode.dirty {
b.Put(keys.RangeTreeNodeKey(cachedNode.node.Key), cachedNode.node)
}
}
return nil
}
// loadNodes fetches a node and recursively all of its children.
func loadNodes(t *testing.T, db *client.DB, key roachpb.RKey, nodes map[string]roachpb.RangeTreeNode) {
node := new(roachpb.RangeTreeNode)
if err := db.GetProto(keys.RangeTreeNodeKey(key), node); err != nil {
t.Fatal(err)
}
nodes[node.Key.String()] = *node
if node.LeftKey != nil {
loadNodes(t, db, node.LeftKey, nodes)
}
if node.RightKey != nil {
loadNodes(t, db, node.RightKey, nodes)
}
}
// SetupRangeTree creates a new RangeTree. This should only be called as part
// of store.BootstrapRange.
func SetupRangeTree(batch engine.Engine, ms *engine.MVCCStats, timestamp roachpb.Timestamp, startKey roachpb.RKey) error {
tree := &RangeTree{
RootKey: startKey,
}
node := &RangeTreeNode{
Key: startKey,
Black: true,
}
if err := engine.MVCCPutProto(batch, ms, keys.RangeTreeRoot, timestamp, nil, tree); err != nil {
return err
}
if err := engine.MVCCPutProto(batch, ms, keys.RangeTreeNodeKey(startKey), timestamp, nil, node); err != nil {
return err
}
return nil
}
// getNode returns the RangeTreeNode for the given key. If the key is nil, nil
// is returned.
func (tc *treeContext) getNode(key roachpb.RKey) (*RangeTreeNode, *roachpb.Error) {
if key == nil {
return nil, nil
}
// First check to see if we have the node cached.
keyString := string(key)
cached, ok := tc.nodes[keyString]
if ok {
return cached.node, nil
}
// We don't have it cached so fetch it and add it to the cache.
node := new(RangeTreeNode)
if pErr := tc.txn.GetProto(keys.RangeTreeNodeKey(key), node); pErr != nil {
return nil, pErr
}
tc.nodes[keyString] = cachedNode{
node: node,
dirty: false,
}
return node, nil
}
// getNode returns the RangeTreeNode for the given key. If the key is nil, nil
// is returned.
func (tc *treeContext) getNode(key *proto.Key) (*proto.RangeTreeNode, error) {
if key == nil {
return nil, nil
}
// First check to see if we have the node cached.
keyString := string(*key)
cached, ok := tc.nodes[keyString]
if ok {
return cached.node, nil
}
// We don't have it cached so fetch it and add it to the cache.
node := &proto.RangeTreeNode{}
if err := tc.txn.GetProto(keys.RangeTreeNodeKey(*key), node); err != nil {
return nil, err
}
tc.nodes[keyString] = cachedNode{
node: node,
dirty: false,
}
return node, nil
}
// treeNodesEqual compares the expectedTree from the provided key to the actual
// nodes retrieved from the db. It recursively calls itself on both left and
// right children if they exist.
func treeNodesEqual(db *client.DB, expected testRangeTree, key proto.Key) error {
expectedNode, ok := expected.Nodes[string(key)]
if !ok {
return util.Errorf("Expected does not contain a node for %s", key)
}
actualNode := &proto.RangeTreeNode{}
if err := db.GetProto(keys.RangeTreeNodeKey(key), actualNode); err != nil {
return err
}
if err := nodesEqual(key, expectedNode, *actualNode); err != nil {
return err
}
if expectedNode.LeftKey != nil {
if err := treeNodesEqual(db, expected, *expectedNode.LeftKey); err != nil {
return err
}
}
if expectedNode.RightKey != nil {
if err := treeNodesEqual(db, expected, *expectedNode.RightKey); err != nil {
return err
}
}
return nil
}
// TestTruncateWithSpanAndDescriptor verifies that a batch request is truncated with a
// range span and the range of a descriptor found in cache.
func TestTruncateWithSpanAndDescriptor(t *testing.T) {
defer leaktest.AfterTest(t)
g, s := makeTestGossip(t)
defer s()
g.SetNodeID(1)
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. When a
// range descriptor is looked up by key "b", return the second
// descriptor whose range is ["a", "c") and partially overlaps
// with the first descriptor's range.
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("a"),
EndKey: roachpb.RKey("c"),
Replicas: []roachpb.ReplicaDescriptor{
{
NodeID: 1,
StoreID: 1,
},
},
}
descDB := mockRangeDescriptorDB(func(key roachpb.RKey, _, _ bool) ([]roachpb.RangeDescriptor, *roachpb.Error) {
desc := descriptor1
if key.Equal(roachpb.RKey("b")) {
desc = descriptor2
}
return []roachpb.RangeDescriptor{desc}, nil
})
// Define our rpcSend stub which checks the span of the batch
// requests. The first request should be the point request on
// "a". The second request should be on "b".
first := true
var testFn rpcSendFn = func(_ rpc.Options, method string, addrs []net.Addr, getArgs func(addr net.Addr) proto.Message, getReply func() proto.Message, _ *rpc.Context) ([]proto.Message, error) {
if method != "Node.Batch" {
return nil, util.Errorf("unexpected method %v", method)
}
ba := getArgs(testAddress).(*roachpb.BatchRequest)
rs := keys.Range(*ba)
if first {
if !(rs.Key.Equal(roachpb.RKey("a")) && rs.EndKey.Equal(roachpb.RKey("a").Next())) {
t.Errorf("Unexpected span [%s,%s)", rs.Key, rs.EndKey)
}
first = false
} else {
if !(rs.Key.Equal(roachpb.RKey("b")) && rs.EndKey.Equal(roachpb.RKey("b").Next())) {
t.Errorf("Unexpected span [%s,%s)", rs.Key, rs.EndKey)
}
}
batchReply := getReply().(*roachpb.BatchResponse)
reply := &roachpb.PutResponse{}
batchReply.Add(reply)
return []proto.Message{batchReply}, nil
}
ctx := &DistSenderContext{
RPCSend: testFn,
RangeDescriptorDB: descDB,
}
ds := NewDistSender(ctx, g)
// Send a batch request contains two puts. In the first
// attempt, the range of the descriptor found in the cache is
// ["a", "b"). The request is truncated to contain only the put
// on "a".
//
// In the second attempt, The range of the descriptor found in
// the cache is ["a", c"), but the put on "a" will not be
// resent. The request is truncated to contain only the put on "b".
ba := roachpb.BatchRequest{}
ba.Txn = &roachpb.Transaction{Name: "test"}
val := roachpb.MakeValueFromString("val")
ba.Add(roachpb.NewPut(keys.RangeTreeNodeKey(roachpb.RKey("a")), val).(*roachpb.PutRequest))
ba.Add(roachpb.NewPut(keys.RangeTreeNodeKey(roachpb.RKey("b")), val).(*roachpb.PutRequest))
_, pErr := ds.Send(context.Background(), ba)
if err := pErr.GoError(); err != nil {
t.Fatal(err)
}
}