func getKeyspaceShards(ctx context.Context, topoServ topo.SrvTopoServer, cell, keyspace string, tabletType topodatapb.TabletType) (string, *topodatapb.SrvKeyspace, []*topodatapb.ShardReference, error) { srvKeyspace, err := topoServ.GetSrvKeyspace(ctx, cell, keyspace) if err != nil { return "", nil, nil, vterrors.NewVitessError( vtrpcpb.ErrorCode_INTERNAL_ERROR, err, "keyspace %v fetch error: %v", keyspace, err, ) } // check if the keyspace has been redirected for this tabletType. for _, sf := range srvKeyspace.ServedFrom { if sf.TabletType == tabletType { keyspace = sf.Keyspace srvKeyspace, err = topoServ.GetSrvKeyspace(ctx, cell, keyspace) if err != nil { return "", nil, nil, vterrors.NewVitessError( vtrpcpb.ErrorCode_INTERNAL_ERROR, err, "keyspace %v fetch error: %v", keyspace, err, ) } } } partition := topoproto.SrvKeyspaceGetPartition(srvKeyspace, tabletType) if partition == nil { return "", nil, nil, vterrors.NewVitessError( vtrpcpb.ErrorCode_INTERNAL_ERROR, err, "No partition found for tabletType %v in keyspace %v", strings.ToLower(tabletType.String()), keyspace, ) } return keyspace, srvKeyspace, partition.ShardReferences, nil }
// findAllKeyspaceShards goes through all serving shards in the topology func findAllKeyspaceShards(ctx context.Context, ts topo.SrvTopoServer, cell string) (map[keyspaceShard]bool, error) { ksNames, err := ts.GetSrvKeyspaceNames(ctx, cell) if err != nil { return nil, err } keyspaceShards := make(map[keyspaceShard]bool) var wg sync.WaitGroup var mu sync.Mutex var errRecorder concurrency.AllErrorRecorder for _, ksName := range ksNames { wg.Add(1) go func(keyspace string) { defer wg.Done() // get SrvKeyspace for cell/keyspace ks, err := ts.GetSrvKeyspace(ctx, cell, keyspace) if err != nil { errRecorder.RecordError(err) return } // get all shard names that are used for serving mu.Lock() for _, ksPartition := range ks.Partitions { for _, shard := range ksPartition.ShardReferences { keyspaceShards[keyspaceShard{ keyspace: keyspace, shard: shard.Name, }] = true } } mu.Unlock() }(ksName) } wg.Wait() if errRecorder.HasErrors() { return nil, errRecorder.Error() } return keyspaceShards, nil }