gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/pickfirst.go (about)

     1  /*
     2   *
     3   * Copyright 2017 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 grpc
    20  
    21  import (
    22  	"errors"
    23  	"fmt"
    24  
    25  	"gitee.com/ks-custle/core-gm/grpc/balancer"
    26  	"gitee.com/ks-custle/core-gm/grpc/connectivity"
    27  )
    28  
    29  // PickFirstBalancerName is the name of the pick_first balancer.
    30  const PickFirstBalancerName = "pick_first"
    31  
    32  func newPickfirstBuilder() balancer.Builder {
    33  	return &pickfirstBuilder{}
    34  }
    35  
    36  type pickfirstBuilder struct{}
    37  
    38  func (*pickfirstBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
    39  	return &pickfirstBalancer{cc: cc}
    40  }
    41  
    42  func (*pickfirstBuilder) Name() string {
    43  	return PickFirstBalancerName
    44  }
    45  
    46  type pickfirstBalancer struct {
    47  	state connectivity.State
    48  	cc    balancer.ClientConn
    49  	sc    balancer.SubConn
    50  }
    51  
    52  func (b *pickfirstBalancer) ResolverError(err error) {
    53  	switch b.state {
    54  	case connectivity.TransientFailure, connectivity.Idle, connectivity.Connecting:
    55  		// Set a failing picker if we don't have a good picker.
    56  		b.cc.UpdateState(balancer.State{ConnectivityState: connectivity.TransientFailure,
    57  			Picker: &picker{err: fmt.Errorf("name resolver error: %v", err)},
    58  		})
    59  	}
    60  	if logger.V(2) {
    61  		logger.Infof("pickfirstBalancer: ResolverError called with error %v", err)
    62  	}
    63  }
    64  
    65  func (b *pickfirstBalancer) UpdateClientConnState(cs balancer.ClientConnState) error {
    66  	if len(cs.ResolverState.Addresses) == 0 {
    67  		b.ResolverError(errors.New("produced zero addresses"))
    68  		return balancer.ErrBadResolverState
    69  	}
    70  	if b.sc == nil {
    71  		var err error
    72  		b.sc, err = b.cc.NewSubConn(cs.ResolverState.Addresses, balancer.NewSubConnOptions{})
    73  		if err != nil {
    74  			if logger.V(2) {
    75  				logger.Errorf("pickfirstBalancer: failed to NewSubConn: %v", err)
    76  			}
    77  			b.state = connectivity.TransientFailure
    78  			b.cc.UpdateState(balancer.State{ConnectivityState: connectivity.TransientFailure,
    79  				Picker: &picker{err: fmt.Errorf("error creating connection: %v", err)},
    80  			})
    81  			return balancer.ErrBadResolverState
    82  		}
    83  		b.state = connectivity.Idle
    84  		b.cc.UpdateState(balancer.State{ConnectivityState: connectivity.Idle, Picker: &picker{result: balancer.PickResult{SubConn: b.sc}}})
    85  		b.sc.Connect()
    86  	} else {
    87  		b.cc.UpdateAddresses(b.sc, cs.ResolverState.Addresses)
    88  		b.sc.Connect()
    89  	}
    90  	return nil
    91  }
    92  
    93  func (b *pickfirstBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
    94  	if logger.V(2) {
    95  		logger.Infof("pickfirstBalancer: UpdateSubConnState: %p, %v", sc, s)
    96  	}
    97  	if b.sc != sc {
    98  		if logger.V(2) {
    99  			logger.Infof("pickfirstBalancer: ignored state change because sc is not recognized")
   100  		}
   101  		return
   102  	}
   103  	b.state = s.ConnectivityState
   104  	if s.ConnectivityState == connectivity.Shutdown {
   105  		b.sc = nil
   106  		return
   107  	}
   108  
   109  	switch s.ConnectivityState {
   110  	case connectivity.Ready:
   111  		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{result: balancer.PickResult{SubConn: sc}}})
   112  	case connectivity.Connecting:
   113  		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{err: balancer.ErrNoSubConnAvailable}})
   114  	case connectivity.Idle:
   115  		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &idlePicker{sc: sc}})
   116  	case connectivity.TransientFailure:
   117  		b.cc.UpdateState(balancer.State{
   118  			ConnectivityState: s.ConnectivityState,
   119  			Picker:            &picker{err: s.ConnectionError},
   120  		})
   121  	}
   122  }
   123  
   124  func (b *pickfirstBalancer) Close() {
   125  }
   126  
   127  func (b *pickfirstBalancer) ExitIdle() {
   128  	if b.sc != nil && b.state == connectivity.Idle {
   129  		b.sc.Connect()
   130  	}
   131  }
   132  
   133  type picker struct {
   134  	result balancer.PickResult
   135  	err    error
   136  }
   137  
   138  func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
   139  	return p.result, p.err
   140  }
   141  
   142  // idlePicker is used when the SubConn is IDLE and kicks the SubConn into
   143  // CONNECTING when Pick is called.
   144  type idlePicker struct {
   145  	sc balancer.SubConn
   146  }
   147  
   148  func (i *idlePicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
   149  	i.sc.Connect()
   150  	return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
   151  }
   152  
   153  func init() {
   154  	balancer.Register(newPickfirstBuilder())
   155  }