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