// TestSendRPCRetry verifies that sendRPC failed on first address but succeed on
// second address, the second reply should be successfully returned back.
func TestSendRPCRetry(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)
}
// Fill RangeDescriptor with 2 replicas
var descriptor = roachpb.RangeDescriptor{
RangeID: 1,
StartKey: roachpb.RKey("a"),
EndKey: roachpb.RKey("z"),
}
for i := 1; i <= 2; i++ {
addr := util.MakeUnresolvedAddr("tcp", fmt.Sprintf("node%d", i))
nd := &roachpb.NodeDescriptor{
NodeID: roachpb.NodeID(i),
Address: util.MakeUnresolvedAddr(addr.Network(), addr.String()),
}
if err := g.AddInfoProto(gossip.MakeNodeIDKey(roachpb.NodeID(i)), nd, time.Hour); err != nil {
t.Fatal(err)
}
descriptor.Replicas = append(descriptor.Replicas, roachpb.ReplicaDescriptor{
NodeID: roachpb.NodeID(i),
StoreID: roachpb.StoreID(i),
})
}
// Define our rpcSend stub which returns success on the second address.
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" {
// reply from first address failed
_ = getReply()
// reply from second address succeed
batchReply := getReply().(*roachpb.BatchResponse)
reply := &roachpb.ScanResponse{}
batchReply.Add(reply)
reply.Rows = append([]roachpb.KeyValue{}, roachpb.KeyValue{Key: roachpb.Key("b"), Value: roachpb.Value{}})
return []proto.Message{batchReply}, nil
}
return nil, util.Errorf("unexpected method %v", method)
}
ctx := &DistSenderContext{
RPCSend: testFn,
RangeDescriptorDB: mockRangeDescriptorDB(func(_ roachpb.RKey, _, _ bool) ([]roachpb.RangeDescriptor, *roachpb.Error) {
return []roachpb.RangeDescriptor{descriptor}, nil
}),
}
ds := NewDistSender(ctx, g)
scan := roachpb.NewScan(roachpb.Key("a"), roachpb.Key("d"), 1)
sr, err := client.SendWrapped(ds, nil, scan)
if err != nil {
t.Fatal(err)
}
if l := len(sr.(*roachpb.ScanResponse).Rows); l != 1 {
t.Fatalf("expected 1 row; got %d", l)
}
}
// TestSendRPCRetry verifies that sendRPC failed on first address but succeed on
// second address, the second reply should be successfully returned back.
func TestSendRPCRetry(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)
}
// Fill RangeDescriptor with 2 replicas
var descriptor = roachpb.RangeDescriptor{
RangeID: 1,
StartKey: roachpb.RKey("a"),
EndKey: roachpb.RKey("z"),
}
for i := 1; i <= 2; i++ {
addr := util.MakeUnresolvedAddr("tcp", fmt.Sprintf("node%d", i))
nd := &roachpb.NodeDescriptor{
NodeID: roachpb.NodeID(i),
Address: util.MakeUnresolvedAddr(addr.Network(), addr.String()),
}
if err := g.AddInfoProto(gossip.MakeNodeIDKey(roachpb.NodeID(i)), nd, time.Hour); err != nil {
t.Fatal(err)
}
descriptor.Replicas = append(descriptor.Replicas, roachpb.ReplicaDescriptor{
NodeID: roachpb.NodeID(i),
StoreID: roachpb.StoreID(i),
})
}
var testFn rpcSendFn = func(_ SendOptions, _ ReplicaSlice,
args roachpb.BatchRequest, _ *rpc.Context) (proto.Message, error) {
batchReply := &roachpb.BatchResponse{}
reply := &roachpb.ScanResponse{}
batchReply.Add(reply)
reply.Rows = append([]roachpb.KeyValue{}, roachpb.KeyValue{Key: roachpb.Key("b"), Value: roachpb.Value{}})
return batchReply, nil
}
ctx := &DistSenderContext{
RPCSend: testFn,
RangeDescriptorDB: mockRangeDescriptorDB(func(_ roachpb.RKey, _, _ bool) ([]roachpb.RangeDescriptor, *roachpb.Error) {
return []roachpb.RangeDescriptor{descriptor}, nil
}),
}
ds := NewDistSender(ctx, g)
scan := roachpb.NewScan(roachpb.Key("a"), roachpb.Key("d"), 1)
sr, err := client.SendWrapped(ds, nil, scan)
if err != nil {
t.Fatal(err)
}
if l := len(sr.(*roachpb.ScanResponse).Rows); l != 1 {
t.Fatalf("expected 1 row; got %d", l)
}
}
// TestSendRPCOrder verifies that sendRPC correctly takes into account the
// leader, attributes and required consistency to determine where to send
// remote requests.
func TestSendRPCOrder(t *testing.T) {
defer leaktest.AfterTest(t)
g, s := makeTestGossip(t)
defer s()
rangeID := roachpb.RangeID(99)
nodeAttrs := map[int32][]string{
1: {}, // The local node, set in each test case.
2: {"us", "west", "gpu"},
3: {"eu", "dublin", "pdu2", "gpu"},
4: {"us", "east", "gpu"},
5: {"us", "east", "gpu", "flaky"},
}
// Gets filled below to identify the replica by its address.
addrToNode := make(map[string]int32)
makeVerifier := func(expOrder rpc.OrderingPolicy,
expAddrs []int32) func(rpc.Options, []net.Addr) error {
return func(o rpc.Options, addrs []net.Addr) error {
if o.Ordering != expOrder {
return util.Errorf("unexpected ordering, wanted %v, got %v",
expOrder, o.Ordering)
}
var actualAddrs []int32
for i, a := range addrs {
if len(expAddrs) <= i {
return util.Errorf("got unexpected address: %s", a)
}
if expAddrs[i] == 0 {
actualAddrs = append(actualAddrs, 0)
} else {
actualAddrs = append(actualAddrs, addrToNode[a.String()])
}
}
if !reflect.DeepEqual(expAddrs, actualAddrs) {
return util.Errorf("expected %d, but found %d", expAddrs, actualAddrs)
}
return nil
}
}
testCases := []struct {
args roachpb.Request
attrs []string
order rpc.OrderingPolicy
expReplica []int32
leader int32 // 0 for not caching a leader.
// Naming is somewhat off, as eventually consistent reads usually
// do not have to go to the leader when a node has a read lease.
// Would really want CONSENSUS here, but that is not implemented.
// Likely a test setup here will never have a read lease, but good
// to keep in mind.
consistent bool
}{
// Inconsistent Scan without matching attributes.
{
args: &roachpb.ScanRequest{},
attrs: []string{},
order: rpc.OrderRandom,
expReplica: []int32{1, 2, 3, 4, 5},
},
// Inconsistent Scan with matching attributes.
// Should move the two nodes matching the attributes to the front and
// go stable.
{
args: &roachpb.ScanRequest{},
attrs: nodeAttrs[5],
order: rpc.OrderStable,
// Compare only the first two resulting addresses.
expReplica: []int32{5, 4, 0, 0, 0},
},
// Scan without matching attributes that requires but does not find
// a leader.
{
args: &roachpb.ScanRequest{},
attrs: []string{},
order: rpc.OrderRandom,
expReplica: []int32{1, 2, 3, 4, 5},
consistent: true,
},
// Put without matching attributes that requires but does not find leader.
// Should go random and not change anything.
{
args: &roachpb.PutRequest{},
attrs: []string{"nomatch"},
order: rpc.OrderRandom,
expReplica: []int32{1, 2, 3, 4, 5},
},
// Put with matching attributes but no leader.
// Should move the two nodes matching the attributes to the front and
// go stable.
{
args: &roachpb.PutRequest{},
attrs: append(nodeAttrs[5], "irrelevant"),
// Compare only the first two resulting addresses.
order: rpc.OrderStable,
expReplica: []int32{5, 4, 0, 0, 0},
},
// Put with matching attributes that finds the leader (node 3).
// Should address the leader and the two nodes matching the attributes
// (the last and second to last) in that order.
{
args: &roachpb.PutRequest{},
attrs: append(nodeAttrs[5], "irrelevant"),
// Compare only the first resulting addresses as we have a leader
// and that means we're only trying to send there.
order: rpc.OrderStable,
expReplica: []int32{2, 5, 4, 0, 0},
leader: 2,
},
// Inconsistent Get without matching attributes but leader (node 3). Should just
// go random as the leader does not matter.
{
args: &roachpb.GetRequest{},
attrs: []string{},
order: rpc.OrderRandom,
expReplica: []int32{1, 2, 3, 4, 5},
leader: 2,
},
}
descriptor := roachpb.RangeDescriptor{
StartKey: roachpb.RKeyMin,
EndKey: roachpb.RKeyMax,
RangeID: rangeID,
Replicas: nil,
}
// Stub to be changed in each test case.
var verifyCall func(rpc.Options, []net.Addr) error
var testFn rpcSendFn = func(opts rpc.Options, method string,
addrs []net.Addr, getArgs func(addr net.Addr) proto.Message,
getReply func() proto.Message, _ *rpc.Context) ([]proto.Message, error) {
if err := verifyCall(opts, addrs); err != nil {
return nil, err
}
return []proto.Message{getArgs(addrs[0]).(*roachpb.BatchRequest).CreateReply()}, nil
}
ctx := &DistSenderContext{
RPCSend: testFn,
RangeDescriptorDB: mockRangeDescriptorDB(func(roachpb.RKey, bool, bool) ([]roachpb.RangeDescriptor, *roachpb.Error) {
return []roachpb.RangeDescriptor{descriptor}, nil
}),
}
ds := NewDistSender(ctx, g)
for n, tc := range testCases {
verifyCall = makeVerifier(tc.order, tc.expReplica)
descriptor.Replicas = nil // could do this once above, but more convenient here
for i := int32(1); i <= 5; i++ {
addr := util.MakeUnresolvedAddr("tcp", fmt.Sprintf("node%d", i))
addrToNode[addr.String()] = i
nd := &roachpb.NodeDescriptor{
NodeID: roachpb.NodeID(i),
Address: util.MakeUnresolvedAddr(addr.Network(), addr.String()),
Attrs: roachpb.Attributes{
Attrs: nodeAttrs[i],
},
}
if err := g.AddInfoProto(gossip.MakeNodeIDKey(roachpb.NodeID(i)), nd, time.Hour); err != nil {
t.Fatal(err)
}
descriptor.Replicas = append(descriptor.Replicas, roachpb.ReplicaDescriptor{
NodeID: roachpb.NodeID(i),
StoreID: roachpb.StoreID(i),
})
}
{
// The local node needs to get its attributes during sendRPC.
nd := &roachpb.NodeDescriptor{
NodeID: 6,
Attrs: roachpb.Attributes{
Attrs: tc.attrs,
},
}
g.SetNodeID(nd.NodeID)
if err := g.SetNodeDescriptor(nd); err != nil {
t.Fatal(err)
}
}
ds.leaderCache.Update(roachpb.RangeID(rangeID), roachpb.ReplicaDescriptor{})
if tc.leader > 0 {
ds.leaderCache.Update(roachpb.RangeID(rangeID), descriptor.Replicas[tc.leader-1])
}
args := tc.args
args.Header().Key = roachpb.Key("a")
if roachpb.IsRange(args) {
args.Header().EndKey = roachpb.Key("b")
}
consistency := roachpb.CONSISTENT
if !tc.consistent {
consistency = roachpb.INCONSISTENT
}
// Kill the cached NodeDescriptor, enforcing a lookup from Gossip.
ds.nodeDescriptor = nil
if _, err := client.SendWrappedWith(ds, nil, roachpb.Header{
RangeID: rangeID, // Not used in this test, but why not.
ReadConsistency: consistency,
}, args); err != nil {
t.Errorf("%d: %s", n, err)
}
}
}