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