gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/balancer/roundrobin/roundrobin.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 roundrobin defines a roundrobin balancer. Roundrobin balancer is
    20  // installed as one of the default balancers in gRPC, users don't need to
    21  // explicitly install this balancer.
    22  package roundrobin
    23  
    24  import (
    25  	"sync"
    26  
    27  	"gitee.com/ks-custle/core-gm/grpc/balancer"
    28  	"gitee.com/ks-custle/core-gm/grpc/balancer/base"
    29  	"gitee.com/ks-custle/core-gm/grpc/grpclog"
    30  	"gitee.com/ks-custle/core-gm/grpc/internal/grpcrand"
    31  )
    32  
    33  // Name is the name of round_robin balancer.
    34  const Name = "round_robin"
    35  
    36  var logger = grpclog.Component("roundrobin")
    37  
    38  // newBuilder creates a new roundrobin balancer builder.
    39  func newBuilder() balancer.Builder {
    40  	return base.NewBalancerBuilder(Name, &rrPickerBuilder{}, base.Config{HealthCheck: true})
    41  }
    42  
    43  func init() {
    44  	balancer.Register(newBuilder())
    45  }
    46  
    47  type rrPickerBuilder struct{}
    48  
    49  func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
    50  	logger.Infof("roundrobinPicker: Build called with info: %v", info)
    51  	if len(info.ReadySCs) == 0 {
    52  		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
    53  	}
    54  	scs := make([]balancer.SubConn, 0, len(info.ReadySCs))
    55  	for sc := range info.ReadySCs {
    56  		scs = append(scs, sc)
    57  	}
    58  	return &rrPicker{
    59  		subConns: scs,
    60  		// Start at a random index, as the same RR balancer rebuilds a new
    61  		// picker when SubConn states change, and we don't want to apply excess
    62  		// load to the first server in the list.
    63  		next: grpcrand.Intn(len(scs)),
    64  	}
    65  }
    66  
    67  type rrPicker struct {
    68  	// subConns is the snapshot of the roundrobin balancer when this picker was
    69  	// created. The slice is immutable. Each Get() will do a round robin
    70  	// selection from it and return the selected SubConn.
    71  	subConns []balancer.SubConn
    72  
    73  	mu   sync.Mutex
    74  	next int
    75  }
    76  
    77  func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
    78  	p.mu.Lock()
    79  	sc := p.subConns[p.next]
    80  	p.next = (p.next + 1) % len(p.subConns)
    81  	p.mu.Unlock()
    82  	return balancer.PickResult{SubConn: sc}, nil
    83  }