google.golang.org/grpc@v1.74.2/balancer/ringhash/picker.go (about)

     1  /*
     2   *
     3   * Copyright 2021 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 ringhash
    20  
    21  import (
    22  	"fmt"
    23  	"strings"
    24  
    25  	xxhash "github.com/cespare/xxhash/v2"
    26  
    27  	"google.golang.org/grpc/balancer"
    28  	"google.golang.org/grpc/connectivity"
    29  	iringhash "google.golang.org/grpc/internal/ringhash"
    30  	"google.golang.org/grpc/metadata"
    31  )
    32  
    33  type picker struct {
    34  	ring *ring
    35  
    36  	// endpointStates is a cache of endpoint states.
    37  	// The ringhash balancer stores endpoint states in a `resolver.EndpointMap`,
    38  	// with access guarded by `ringhashBalancer.mu`. The `endpointStates` cache
    39  	// in the picker helps avoid locking the ringhash balancer's mutex when
    40  	// reading the latest state at RPC time.
    41  	endpointStates map[string]endpointState // endpointState.hashKey -> endpointState
    42  
    43  	// requestHashHeader is the header key to look for the request hash. If it's
    44  	// empty, the request hash is expected to be set in the context via xDS.
    45  	// See gRFC A76.
    46  	requestHashHeader string
    47  
    48  	// hasEndpointInConnectingState is true if any of the endpoints is in
    49  	// CONNECTING.
    50  	hasEndpointInConnectingState bool
    51  
    52  	randUint64 func() uint64
    53  }
    54  
    55  func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
    56  	usingRandomHash := false
    57  	var requestHash uint64
    58  	if p.requestHashHeader == "" {
    59  		var ok bool
    60  		if requestHash, ok = iringhash.XDSRequestHash(info.Ctx); !ok {
    61  			return balancer.PickResult{}, fmt.Errorf("ringhash: expected xDS config selector to set the request hash")
    62  		}
    63  	} else {
    64  		md, ok := metadata.FromOutgoingContext(info.Ctx)
    65  		if !ok || len(md.Get(p.requestHashHeader)) == 0 {
    66  			requestHash = p.randUint64()
    67  			usingRandomHash = true
    68  		} else {
    69  			values := strings.Join(md.Get(p.requestHashHeader), ",")
    70  			requestHash = xxhash.Sum64String(values)
    71  		}
    72  	}
    73  
    74  	e := p.ring.pick(requestHash)
    75  	ringSize := len(p.ring.items)
    76  	if !usingRandomHash {
    77  		// Per gRFC A61, because of sticky-TF with PickFirst's auto reconnect on TF,
    78  		// we ignore all TF subchannels and find the first ring entry in READY,
    79  		// CONNECTING or IDLE.  If that entry is in IDLE, we need to initiate a
    80  		// connection. The idlePicker returned by the LazyLB or the new Pickfirst
    81  		// should do this automatically.
    82  		for i := 0; i < ringSize; i++ {
    83  			index := (e.idx + i) % ringSize
    84  			es := p.endpointState(p.ring.items[index])
    85  			switch es.state.ConnectivityState {
    86  			case connectivity.Ready, connectivity.Connecting, connectivity.Idle:
    87  				return es.state.Picker.Pick(info)
    88  			case connectivity.TransientFailure:
    89  			default:
    90  				panic(fmt.Sprintf("Found child balancer in unknown state: %v", es.state.ConnectivityState))
    91  			}
    92  		}
    93  	} else {
    94  		// If the picker has generated a random hash, it will walk the ring from
    95  		// this hash, and pick the first READY endpoint. If no endpoint is
    96  		// currently in CONNECTING state, it will trigger a connection attempt
    97  		// on at most one endpoint that is in IDLE state along the way. - A76
    98  		requestedConnection := p.hasEndpointInConnectingState
    99  		for i := 0; i < ringSize; i++ {
   100  			index := (e.idx + i) % ringSize
   101  			es := p.endpointState(p.ring.items[index])
   102  			if es.state.ConnectivityState == connectivity.Ready {
   103  				return es.state.Picker.Pick(info)
   104  			}
   105  			if !requestedConnection && es.state.ConnectivityState == connectivity.Idle {
   106  				requestedConnection = true
   107  				// If the SubChannel is in idle state, initiate a connection but
   108  				// continue to check other pickers to see if there is one in
   109  				// ready state.
   110  				es.balancer.ExitIdle()
   111  			}
   112  		}
   113  		if requestedConnection {
   114  			return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
   115  		}
   116  	}
   117  
   118  	// All children are in transient failure. Return the first failure.
   119  	return p.endpointState(e).state.Picker.Pick(info)
   120  }
   121  
   122  func (p *picker) endpointState(e *ringEntry) endpointState {
   123  	return p.endpointStates[e.hashKey]
   124  }