vendor/google.golang.org/grpc/balancer/roundrobin/roundrobin.go

   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         "context"
  26         "sync"
  27
  28         "google.golang.org/grpc/balancer"
  29         "google.golang.org/grpc/balancer/base"
  30         "google.golang.org/grpc/grpclog"
  31         "google.golang.org/grpc/resolver"
  32 )
  33
  34 // Name is the name of round_robin balancer.
  35 const Name = "round_robin"
  36
  37 // newBuilder creates a new roundrobin balancer builder.
  38 func newBuilder() balancer.Builder {
  39         return base.NewBalancerBuilderWithConfig(Name, &rrPickerBuilder{}, base.Config{HealthCheck: true})
  40 }
  41
  42 func init() {
  43         balancer.Register(newBuilder())
  44 }
  45
  46 type rrPickerBuilder struct{}
  47
  48 func (*rrPickerBuilder) Build(readySCs map[resolver.Address]balancer.SubConn) balancer.Picker {
  49         grpclog.Infof("roundrobinPicker: newPicker called with readySCs: %v", readySCs)
  50         var scs []balancer.SubConn
  51         for _, sc := range readySCs {
  52                 scs = append(scs, sc)
  53         }
  54         return &rrPicker{
  55                 subConns: scs,
  56         }
  57 }
  58
  59 type rrPicker struct {
  60         // subConns is the snapshot of the roundrobin balancer when this picker was
  61         // created. The slice is immutable. Each Get() will do a round robin
  62         // selection from it and return the selected SubConn.
  63         subConns []balancer.SubConn
  64
  65         mu   sync.Mutex
  66         next int
  67 }
  68
  69 func (p *rrPicker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
  70         if len(p.subConns) <= 0 {
  71                 return nil, nil, balancer.ErrNoSubConnAvailable
  72         }
  73
  74         p.mu.Lock()
  75         sc := p.subConns[p.next]
  76         p.next = (p.next + 1) % len(p.subConns)
  77         p.mu.Unlock()
  78         return sc, nil, nil
  79 }