dubbo.apache.org/dubbo-go/v3@v3.1.1/xds/balancer/clustermanager/picker.go (about)

     1  /*
     2   * Licensed to the Apache Software Foundation (ASF) under one or more
     3   * contributor license agreements.  See the NOTICE file distributed with
     4   * this work for additional information regarding copyright ownership.
     5   * The ASF licenses this file to You under the Apache License, Version 2.0
     6   * (the "License"); you may not use this file except in compliance with
     7   * the License.  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   *
    20   * Copyright 2020 gRPC authors.
    21   *
    22   */
    23  
    24  package clustermanager
    25  
    26  import (
    27  	"context"
    28  )
    29  
    30  import (
    31  	"google.golang.org/grpc/balancer"
    32  
    33  	"google.golang.org/grpc/codes"
    34  
    35  	"google.golang.org/grpc/status"
    36  )
    37  
    38  // pickerGroup contains a list of pickers. If the picker isn't ready, the pick
    39  // will be queued.
    40  type pickerGroup struct {
    41  	pickers map[string]balancer.Picker
    42  }
    43  
    44  func newPickerGroup(idToPickerState map[string]*subBalancerState) *pickerGroup {
    45  	pickers := make(map[string]balancer.Picker)
    46  	for id, st := range idToPickerState {
    47  		pickers[id] = st.state.Picker
    48  	}
    49  	return &pickerGroup{
    50  		pickers: pickers,
    51  	}
    52  }
    53  
    54  func (pg *pickerGroup) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
    55  	cluster := getPickedCluster(info.Ctx)
    56  	if p := pg.pickers[cluster]; p != nil {
    57  		return p.Pick(info)
    58  	}
    59  	return balancer.PickResult{}, status.Errorf(codes.Unavailable, "unknown cluster selected for RPC: %q", cluster)
    60  }
    61  
    62  type clusterKey struct{}
    63  
    64  func getPickedCluster(ctx context.Context) string {
    65  	cluster, _ := ctx.Value(clusterKey{}).(string)
    66  	return cluster
    67  }
    68  
    69  // GetPickedClusterForTesting returns the cluster in the context; to be used
    70  // for testing only.
    71  func GetPickedClusterForTesting(ctx context.Context) string {
    72  	return getPickedCluster(ctx)
    73  }
    74  
    75  // SetPickedCluster adds the selected cluster to the context for the
    76  // xds_cluster_manager LB policy to pick.
    77  func SetPickedCluster(ctx context.Context, cluster string) context.Context {
    78  	return context.WithValue(ctx, clusterKey{}, cluster)
    79  }