func TestMixedDirections(t *testing.T) {
defer leaktest.AfterTest(t)()
s, db, cdb := serverutils.StartServer(t, base.TestServerArgs{
UseDatabase: "t",
})
defer s.Stopper().Stop()
rowRanges, tableDesc := setupRanges(db, s.(*server.TestServer), cdb, t)
lr := distsqlplan.NewSpanResolver(
s.DistSender(), s.Gossip(),
s.(*server.TestServer).GetNode().Descriptor,
distsqlplan.BinPackingLeaseHolderChoice)
ctx := context.Background()
it := lr.NewSpanResolverIterator()
spans := []spanWithDir{
orient(kv.Ascending, makeSpan(tableDesc, 11, 15))[0],
orient(kv.Descending, makeSpan(tableDesc, 1, 14))[0],
}
replicas, err := resolveSpans(ctx, it, spans...)
if err != nil {
t.Fatal(err)
}
expected := [][]rngInfo{
{onlyReplica(rowRanges[1])},
{onlyReplica(rowRanges[1]), onlyReplica(rowRanges[0])},
}
if err = expectResolved(replicas, expected...); err != nil {
t.Fatal(err)
}
}
func newDistSQLPlanner(
nodeDesc roachpb.NodeDescriptor,
rpcCtx *rpc.Context,
distSQLSrv *distsqlrun.ServerImpl,
distSender *kv.DistSender,
gossip *gossip.Gossip,
) *distSQLPlanner {
return &distSQLPlanner{
nodeDesc: nodeDesc,
rpcContext: rpcCtx,
distSQLSrv: distSQLSrv,
spanResolver: distsqlplan.NewSpanResolver(distSender, gossip, nodeDesc, resolverPolicy),
}
}
// Test that resolving spans uses a node's range cache and lease holder cache.
// The idea is to test that resolving is not random, but predictable given the
// state of caches.
func TestSpanResolverUsesCaches(t *testing.T) {
defer leaktest.AfterTest(t)()
tc := testcluster.StartTestCluster(t, 4,
base.TestClusterArgs{
ReplicationMode: base.ReplicationManual,
ServerArgs: base.TestServerArgs{
UseDatabase: "t",
},
})
defer tc.Stopper().Stop()
rowRanges, _ := setupRanges(
tc.Conns[0], tc.Servers[0], tc.Servers[0].KVClient().(*client.DB), t)
// Replicate the row ranges on all of the first 3 nodes. Save the 4th node in
// a pristine state, with empty caches.
for i := 0; i < 3; i++ {
var err error
rowRanges[i], err = tc.AddReplicas(
rowRanges[i].StartKey.AsRawKey(), tc.Target(1), tc.Target(2))
if err != nil {
t.Fatal(err)
}
}
// Scatter the leases around; node i gets range i.
for i := 0; i < 3; i++ {
if err := tc.TransferRangeLease(rowRanges[i], tc.Target(i)); err != nil {
t.Fatal(err)
}
// Wait for everybody to apply the new lease, so that we can rely on the
// lease discovery done later by the SpanResolver to be up to date.
testutils.SucceedsSoon(t, func() error {
for j := 0; j < 3; j++ {
target := tc.Target(j)
rt, err := tc.FindRangeLeaseHolder(rowRanges[i], &target)
if err != nil {
return err
}
if rt != tc.Target(i) {
return errors.Errorf("node %d hasn't applied the lease yet", j)
}
}
return nil
})
}
// Create a SpanResolver using the 4th node, with empty caches.
s3 := tc.Servers[3]
lr := distsqlplan.NewSpanResolver(
s3.DistSender(), s3.Gossip(), s3.GetNode().Descriptor,
distsqlplan.BinPackingLeaseHolderChoice)
var spans []spanWithDir
for i := 0; i < 3; i++ {
spans = append(
spans,
spanWithDir{
Span: roachpb.Span{
Key: rowRanges[i].StartKey.AsRawKey(),
EndKey: rowRanges[i].EndKey.AsRawKey(),
},
dir: kv.Ascending,
})
}
// Resolve the spans. Since the LeaseHolderCache is empty, all the ranges
// should be grouped and "assigned" to replica 0.
replicas, err := resolveSpans(context.TODO(), lr.NewSpanResolverIterator(), spans...)
if err != nil {
t.Fatal(err)
}
if len(replicas) != 3 {
t.Fatalf("expected replies for 3 spans, got %d: %+v", len(replicas), replicas)
}
si := tc.Servers[0]
nodeID := si.GetNode().Descriptor.NodeID
storeID := si.GetFirstStoreID()
for i := 0; i < 3; i++ {
if len(replicas[i]) != 1 {
t.Fatalf("expected 1 range for span %s, got %d (%+v)",
spans[i].Span, len(replicas[i]), replicas[i])
}
rd := replicas[i][0].ReplicaDescriptor
if rd.NodeID != nodeID || rd.StoreID != storeID {
t.Fatalf("expected span %s to be on replica (%d, %d) but was on %s",
spans[i].Span, nodeID, storeID, rd)
}
}
// Now populate the cached on node 4 and query again. This time, we expect to see
// each span on its own range.
if err := populateCache(tc.Conns[3], 3 /* expectedNumRows */); err != nil {
t.Fatal(err)
}
replicas, err = resolveSpans(context.TODO(), lr.NewSpanResolverIterator(), spans...)
if err != nil {
t.Fatal(err)
}
var expected [][]rngInfo
for i := 0; i < 3; i++ {
expected = append(expected, []rngInfo{selectReplica(tc.Servers[i].NodeID(), rowRanges[i])})
}
if err = expectResolved(replicas, expected...); err != nil {
t.Fatal(err)
}
}
func TestSpanResolver(t *testing.T) {
defer leaktest.AfterTest(t)()
s, db, cdb := serverutils.StartServer(t, base.TestServerArgs{
UseDatabase: "t",
})
defer s.Stopper().Stop()
rowRanges, tableDesc := setupRanges(db, s.(*server.TestServer), cdb, t)
lr := distsqlplan.NewSpanResolver(
s.DistSender(), s.Gossip(),
s.(*server.TestServer).GetNode().Descriptor,
distsqlplan.BinPackingLeaseHolderChoice)
ctx := context.Background()
it := lr.NewSpanResolverIterator()
testCases := []struct {
spans []roachpb.Span
expected [][]rngInfo
}{
{
[]roachpb.Span{makeSpan(tableDesc, 0, 10000)},
[][]rngInfo{{
onlyReplica(rowRanges[0]),
onlyReplica(rowRanges[1]),
onlyReplica(rowRanges[2])}},
},
{
[]roachpb.Span{
makeSpan(tableDesc, 0, 9),
makeSpan(tableDesc, 11, 19),
makeSpan(tableDesc, 21, 29),
},
[][]rngInfo{
{onlyReplica(rowRanges[0])},
{onlyReplica(rowRanges[1])},
{onlyReplica(rowRanges[2])},
},
},
{
[]roachpb.Span{
makeSpan(tableDesc, 0, 20),
makeSpan(tableDesc, 20, 29),
},
[][]rngInfo{
{onlyReplica(rowRanges[0]), onlyReplica(rowRanges[1])},
{onlyReplica(rowRanges[2])},
},
},
{
[]roachpb.Span{
makeSpan(tableDesc, 0, 1),
makeSpan(tableDesc, 1, 2),
makeSpan(tableDesc, 2, 3),
makeSpan(tableDesc, 3, 4),
makeSpan(tableDesc, 5, 11),
makeSpan(tableDesc, 20, 29),
},
[][]rngInfo{
{onlyReplica(rowRanges[0])},
{onlyReplica(rowRanges[0])},
{onlyReplica(rowRanges[0])},
{onlyReplica(rowRanges[0])},
{onlyReplica(rowRanges[0]), onlyReplica(rowRanges[1])},
{onlyReplica(rowRanges[2])},
},
},
}
for i, tc := range testCases {
for _, dir := range []kv.ScanDirection{kv.Ascending, kv.Descending} {
t.Run(fmt.Sprintf("%d-direction:%d", i, dir), func(t *testing.T) {
replicas, err := resolveSpans(ctx, it, orient(dir, tc.spans...)...)
if err != nil {
t.Fatal(err)
}
if dir == kv.Descending {
// When testing Descending resolving, reverse the expected results.
for i, j := 0, len(tc.expected)-1; i <= j; i, j = i+1, j-1 {
reverse(tc.expected[i])
if i != j {
reverse(tc.expected[j])
}
tc.expected[i], tc.expected[j] = tc.expected[j], tc.expected[i]
}
}
if err = expectResolved(replicas, tc.expected...); err != nil {
t.Fatal(err)
}
})
}
}
}