1/* 2 * 3 * Copyright 2020 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19package clustermanager 20 21import ( 22 "context" 23 24 "google.golang.org/grpc/balancer" 25 "google.golang.org/grpc/codes" 26 "google.golang.org/grpc/status" 27) 28 29// pickerGroup contains a list of pickers. If the picker isn't ready, the pick 30// will be queued. 31type pickerGroup struct { 32 pickers map[string]balancer.Picker 33} 34 35func newPickerGroup(idToPickerState map[string]*subBalancerState) *pickerGroup { 36 pickers := make(map[string]balancer.Picker) 37 for id, st := range idToPickerState { 38 pickers[id] = st.state.Picker 39 } 40 return &pickerGroup{ 41 pickers: pickers, 42 } 43} 44 45func (pg *pickerGroup) Pick(info balancer.PickInfo) (balancer.PickResult, error) { 46 cluster := getPickedCluster(info.Ctx) 47 if p := pg.pickers[cluster]; p != nil { 48 return p.Pick(info) 49 } 50 return balancer.PickResult{}, status.Errorf(codes.Unavailable, "unknown cluster selected for RPC: %q", cluster) 51} 52 53type clusterKey struct{} 54 55func getPickedCluster(ctx context.Context) string { 56 cluster, _ := ctx.Value(clusterKey{}).(string) 57 return cluster 58} 59 60// GetPickedClusterForTesting returns the cluster in the context; to be used 61// for testing only. 62func GetPickedClusterForTesting(ctx context.Context) string { 63 return getPickedCluster(ctx) 64} 65 66// SetPickedCluster adds the selected cluster to the context for the 67// xds_cluster_manager LB policy to pick. 68func SetPickedCluster(ctx context.Context, cluster string) context.Context { 69 return context.WithValue(ctx, clusterKey{}, cluster) 70} 71