Skip to content

Commit

Permalink
test real tablets result
Browse files Browse the repository at this point in the history 
Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
  • Loading branch information
@timvaillancourt
timvaillancourt committed Nov 5, 2024
1 parent 60ba828 commit 90d249f
Show file tree
Hide file tree
Showing 3 changed files with 238 additions and 39 deletions.
13 changes: 8 additions & 5 deletions go/vt/topo/shard.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -654,12 +654,16 @@ func (ts *Server) GetTabletsByShardCell(ctx context.Context, keyspace, shard str
eg.SetLimit(DefaultConcurrency)

tablets := make([]*TabletInfo, 0, len(cells))
options := &GetTabletsByCellOptions{
Concurrency: cellConcurrency,
KeyspaceShard: &KeyspaceShard{
var kss *KeyspaceShard
if keyspace != "" {
kss = &KeyspaceShard{
Keyspace: keyspace,
Shard: shard,
},
}
}
options := &GetTabletsByCellOptions{
Concurrency: cellConcurrency,
KeyspaceShard: kss,
}
for _, cell := range cells {
eg.Go(func() error {
Expand All @@ -677,7 +681,6 @@ func (ts *Server) GetTabletsByShardCell(ctx context.Context, keyspace, shard str
log.Warningf("GetTabletsByShardCell(%v,%v): got partial result: %v", keyspace, shard, err)
return tablets, NewError(PartialResult, shard)
}

return tablets, nil
}

Expand Down
10 changes: 7 additions & 3 deletions go/vt/topo/tablet.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,11 @@ limitations under the License.
package topo

import (
"cmp"
"context"
"fmt"
"path"
"slices"
"sort"
"sync"
"time"
Expand Down Expand Up @@ -277,8 +279,8 @@ func (ts *Server) GetTabletsByCell(ctx context.Context, cellAlias string, opt *G
if err := tablet.UnmarshalVT(listResults[n].Value); err != nil {
return nil, err
}
if opt != nil && opt.KeyspaceShard != nil {
if opt.KeyspaceShard.Keyspace != "" && opt.KeyspaceShard.Keyspace != tablet.Keyspace {
if opt != nil && opt.KeyspaceShard != nil && opt.KeyspaceShard.Keyspace != "" {
if opt.KeyspaceShard.Keyspace != tablet.Keyspace {
continue
}
if opt.KeyspaceShard.Shard != "" && opt.KeyspaceShard.Shard != tablet.Shard {
Expand All @@ -287,7 +289,9 @@ func (ts *Server) GetTabletsByCell(ctx context.Context, cellAlias string, opt *G
}
tablets = append(tablets, &TabletInfo{Tablet: tablet, version: listResults[n].Version})
}

slices.SortFunc(tablets, func(i, j *TabletInfo) int {
return cmp.Compare(i.Alias.Uid, j.Alias.Uid)
})
return tablets, nil
}

Expand Down
254 changes: 223 additions & 31 deletions go/vt/topo/tablet_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,68 +41,211 @@ func TestServerGetTabletsByCell(t *testing.T) {
tests := []struct {
name string
createShardTablets int
expectedTablets int
expectedShards []string
expectedTablets []*topodatapb.Tablet
opt *topo.GetTabletsByCellOptions
listError error
keyspaceShards map[string][]string
}{
{
name: "negative concurrency",
keyspaceShards: map[string][]string{
keyspace: []string{shard},
keyspace: {shard},
},
createShardTablets: 1,
expectedTablets: 1,
expectedShards: []string{shard},
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
},
// Ensure this doesn't panic.
opt: &topo.GetTabletsByCellOptions{Concurrency: -1},
},
{
name: "single",
keyspaceShards: map[string][]string{
keyspace: []string{shard},
keyspace: {shard},
},
createShardTablets: 1,
expectedTablets: 1,
expectedShards: []string{shard},
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
},
// Make sure the defaults apply as expected.
opt: nil,
},
{
name: "multiple",
keyspaceShards: map[string][]string{
keyspace: []string{shard},
keyspace: {shard},
},
// Should work with more than 1 tablet
createShardTablets: 32,
expectedTablets: 32,
expectedShards: []string{shard},
opt: &topo.GetTabletsByCellOptions{Concurrency: 8},
createShardTablets: 4,
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(2),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(2),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(3),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(3),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(4),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(4),
},
Keyspace: keyspace,
Shard: shard,
},
},
opt: &topo.GetTabletsByCellOptions{Concurrency: 8},
},
{
name: "multiple with list error",
keyspaceShards: map[string][]string{
keyspace: []string{shard},
keyspace: {shard},
},
// Should work with more than 1 tablet when List returns an error
createShardTablets: 32,
expectedTablets: 32,
expectedShards: []string{shard},
createShardTablets: 4,
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(2),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(2),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(3),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(3),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(4),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(4),
},
Keyspace: keyspace,
Shard: shard,
},
},
opt: &topo.GetTabletsByCellOptions{Concurrency: 8},
listError: topo.NewError(topo.ResourceExhausted, ""),
},
{
name: "filtered by keyspace and shard",
keyspaceShards: map[string][]string{
keyspace: []string{shard},
"filtered": []string{"-"},
keyspace: {shard},
"filtered": {"-"},
},
// Should create 2 tablets in 2 different shards (4 total)
// but only a single shard is returned
createShardTablets: 2,
expectedTablets: 2,
expectedShards: []string{shard},
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(2),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(2),
},
Keyspace: keyspace,
Shard: shard,
},
},
opt: &topo.GetTabletsByCellOptions{
Concurrency: 1,
KeyspaceShard: &topo.KeyspaceShard{
Expand All @@ -114,7 +257,7 @@ func TestServerGetTabletsByCell(t *testing.T) {
{
name: "filtered by keyspace and no shard",
keyspaceShards: map[string][]string{
keyspace: []string{
keyspace: {
shard,
shard + "2",
},
Expand All @@ -123,8 +266,56 @@ func TestServerGetTabletsByCell(t *testing.T) {
// in the same keyspace and both shards are returned due to
// empty shard
createShardTablets: 2,
expectedTablets: 4,
expectedShards: []string{shard, shard + "2"},
expectedTablets: []*topodatapb.Tablet{
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(1),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(1),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(2),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(2),
},
Keyspace: keyspace,
Shard: shard,
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(3),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(3),
},
Keyspace: keyspace,
Shard: shard + "2",
},
{
Alias: &topodatapb.TabletAlias{
Cell: cell,
Uid: uint32(4),
},
Hostname: "host1",
PortMap: map[string]int32{
"vt": int32(4),
},
Keyspace: keyspace,
Shard: shard + "2",
},
},
opt: &topo.GetTabletsByCellOptions{
Concurrency: 1,
KeyspaceShard: &topo.KeyspaceShard{
Expand Down Expand Up @@ -155,7 +346,7 @@ func TestServerGetTabletsByCell(t *testing.T) {
}
}

tablets := make([]*topo.TabletInfo, tt.expectedTablets)
tablets := make([]*topo.TabletInfo, 0)

var uid uint32 = 1
for k, shards := range tt.keyspaceShards {
Expand All @@ -174,7 +365,7 @@ func TestServerGetTabletsByCell(t *testing.T) {
Shard: s,
}
tInfo := &topo.TabletInfo{Tablet: tablet}
tablets[i] = tInfo
tablets = append(tablets, tInfo)
require.NoError(t, ts.CreateTablet(ctx, tablet))
uid++
}
Expand All @@ -185,11 +376,12 @@ func TestServerGetTabletsByCell(t *testing.T) {
// tablet matches what we expect.
out, err := ts.GetTabletsByCell(ctx, cell, tt.opt)
require.NoError(t, err)
require.Len(t, out, tt.expectedTablets)

for _, tablet := range out {
require.Equal(t, keyspace, tablet.Keyspace)
require.Contains(t, tt.expectedShards, tablet.Shard)
require.Len(t, out, len(tt.expectedTablets))
for i, tablet := range out {
expected := tt.expectedTablets[i]
require.Equal(t, expected.Alias.String(), tablet.Alias.String())
require.Equal(t, expected.Keyspace, tablet.Keyspace)
require.Equal(t, expected.Shard, tablet.Shard)
}
})
}
Expand Down

0 comments on commit 90d249f

Please sign in to comment.