picker.go

Documentation: google.golang.org/grpc/xds/internal/balancer/clustermanager

     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  
    19  package clustermanager
    20  
    21  import (
    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.
    31  type pickerGroup struct {
    32  	pickers map[string]balancer.Picker
    33  }
    34  
    35  func 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  
    45  func (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  
    53  type clusterKey struct{}
    54  
    55  func 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.
    62  func 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.
    68  func SetPickedCluster(ctx context.Context, cluster string) context.Context {
    69  	return context.WithValue(ctx, clusterKey{}, cluster)
    70  }
    71
View as plain text