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

   1 /*
   2  *
   3  * Copyright 2016 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         "fmt"
  23         "net"
  24         "sync"
  25
  26         "golang.org/x/net/context"
  27         "google.golang.org/grpc/codes"
  28         "google.golang.org/grpc/credentials"
  29         "google.golang.org/grpc/grpclog"
  30         "google.golang.org/grpc/naming"
  31 )
  32
  33 // Address represents a server the client connects to.
  34 // This is the EXPERIMENTAL API and may be changed or extended in the future.
  35 type Address struct {
  36         // Addr is the server address on which a connection will be established.
  37         Addr string
  38         // Metadata is the information associated with Addr, which may be used
  39         // to make load balancing decision.
  40         Metadata interface{}
  41 }
  42
  43 // BalancerConfig specifies the configurations for Balancer.
  44 type BalancerConfig struct {
  45         // DialCreds is the transport credential the Balancer implementation can
  46         // use to dial to a remote load balancer server. The Balancer implementations
  47         // can ignore this if it does not need to talk to another party securely.
  48         DialCreds credentials.TransportCredentials
  49         // Dialer is the custom dialer the Balancer implementation can use to dial
  50         // to a remote load balancer server. The Balancer implementations
  51         // can ignore this if it doesn't need to talk to remote balancer.
  52         Dialer func(context.Context, string) (net.Conn, error)
  53 }
  54
  55 // BalancerGetOptions configures a Get call.
  56 // This is the EXPERIMENTAL API and may be changed or extended in the future.
  57 type BalancerGetOptions struct {
  58         // BlockingWait specifies whether Get should block when there is no
  59         // connected address.
  60         BlockingWait bool
  61 }
  62
  63 // Balancer chooses network addresses for RPCs.
  64 // This is the EXPERIMENTAL API and may be changed or extended in the future.
  65 type Balancer interface {
  66         // Start does the initialization work to bootstrap a Balancer. For example,
  67         // this function may start the name resolution and watch the updates. It will
  68         // be called when dialing.
  69         Start(target string, config BalancerConfig) error
  70         // Up informs the Balancer that gRPC has a connection to the server at
  71         // addr. It returns down which is called once the connection to addr gets
  72         // lost or closed.
  73         // TODO: It is not clear how to construct and take advantage of the meaningful error
  74         // parameter for down. Need realistic demands to guide.
  75         Up(addr Address) (down func(error))
  76         // Get gets the address of a server for the RPC corresponding to ctx.
  77         // i) If it returns a connected address, gRPC internals issues the RPC on the
  78         // connection to this address;
  79         // ii) If it returns an address on which the connection is under construction
  80         // (initiated by Notify(...)) but not connected, gRPC internals
  81         //  * fails RPC if the RPC is fail-fast and connection is in the TransientFailure or
  82         //  Shutdown state;
  83         //  or
  84         //  * issues RPC on the connection otherwise.
  85         // iii) If it returns an address on which the connection does not exist, gRPC
  86         // internals treats it as an error and will fail the corresponding RPC.
  87         //
  88         // Therefore, the following is the recommended rule when writing a custom Balancer.
  89         // If opts.BlockingWait is true, it should return a connected address or
  90         // block if there is no connected address. It should respect the timeout or
  91         // cancellation of ctx when blocking. If opts.BlockingWait is false (for fail-fast
  92         // RPCs), it should return an address it has notified via Notify(...) immediately
  93         // instead of blocking.
  94         //
  95         // The function returns put which is called once the rpc has completed or failed.
  96         // put can collect and report RPC stats to a remote load balancer.
  97         //
  98         // This function should only return the errors Balancer cannot recover by itself.
  99         // gRPC internals will fail the RPC if an error is returned.
 100         Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error)
 101         // Notify returns a channel that is used by gRPC internals to watch the addresses
 102         // gRPC needs to connect. The addresses might be from a name resolver or remote
 103         // load balancer. gRPC internals will compare it with the existing connected
 104         // addresses. If the address Balancer notified is not in the existing connected
 105         // addresses, gRPC starts to connect the address. If an address in the existing
 106         // connected addresses is not in the notification list, the corresponding connection
 107         // is shutdown gracefully. Otherwise, there are no operations to take. Note that
 108         // the Address slice must be the full list of the Addresses which should be connected.
 109         // It is NOT delta.
 110         Notify() <-chan []Address
 111         // Close shuts down the balancer.
 112         Close() error
 113 }
 114
 115 // downErr implements net.Error. It is constructed by gRPC internals and passed to the down
 116 // call of Balancer.
 117 type downErr struct {
 118         timeout   bool
 119         temporary bool
 120         desc      string
 121 }
 122
 123 func (e downErr) Error() string   { return e.desc }
 124 func (e downErr) Timeout() bool   { return e.timeout }
 125 func (e downErr) Temporary() bool { return e.temporary }
 126
 127 func downErrorf(timeout, temporary bool, format string, a ...interface{}) downErr {
 128         return downErr{
 129                 timeout:   timeout,
 130                 temporary: temporary,
 131                 desc:      fmt.Sprintf(format, a...),
 132         }
 133 }
 134
 135 // RoundRobin returns a Balancer that selects addresses round-robin. It uses r to watch
 136 // the name resolution updates and updates the addresses available correspondingly.
 137 func RoundRobin(r naming.Resolver) Balancer {
 138         return &roundRobin{r: r}
 139 }
 140
 141 type addrInfo struct {
 142         addr      Address
 143         connected bool
 144 }
 145
 146 type roundRobin struct {
 147         r      naming.Resolver
 148         w      naming.Watcher
 149         addrs  []*addrInfo // all the addresses the client should potentially connect
 150         mu     sync.Mutex
 151         addrCh chan []Address // the channel to notify gRPC internals the list of addresses the client should connect to.
 152         next   int            // index of the next address to return for Get()
 153         waitCh chan struct{}  // the channel to block when there is no connected address available
 154         done   bool           // The Balancer is closed.
 155 }
 156
 157 func (rr *roundRobin) watchAddrUpdates() error {
 158         updates, err := rr.w.Next()
 159         if err != nil {
 160                 grpclog.Warningf("grpc: the naming watcher stops working due to %v.", err)
 161                 return err
 162         }
 163         rr.mu.Lock()
 164         defer rr.mu.Unlock()
 165         for _, update := range updates {
 166                 addr := Address{
 167                         Addr:     update.Addr,
 168                         Metadata: update.Metadata,
 169                 }
 170                 switch update.Op {
 171                 case naming.Add:
 172                         var exist bool
 173                         for _, v := range rr.addrs {
 174                                 if addr == v.addr {
 175                                         exist = true
 176                                         grpclog.Infoln("grpc: The name resolver wanted to add an existing address: ", addr)
 177                                         break
 178                                 }
 179                         }
 180                         if exist {
 181                                 continue
 182                         }
 183                         rr.addrs = append(rr.addrs, &addrInfo{addr: addr})
 184                 case naming.Delete:
 185                         for i, v := range rr.addrs {
 186                                 if addr == v.addr {
 187                                         copy(rr.addrs[i:], rr.addrs[i+1:])
 188                                         rr.addrs = rr.addrs[:len(rr.addrs)-1]
 189                                         break
 190                                 }
 191                         }
 192                 default:
 193                         grpclog.Errorln("Unknown update.Op ", update.Op)
 194                 }
 195         }
 196         // Make a copy of rr.addrs and write it onto rr.addrCh so that gRPC internals gets notified.
 197         open := make([]Address, len(rr.addrs))
 198         for i, v := range rr.addrs {
 199                 open[i] = v.addr
 200         }
 201         if rr.done {
 202                 return ErrClientConnClosing
 203         }
 204         select {
 205         case <-rr.addrCh:
 206         default:
 207         }
 208         rr.addrCh <- open
 209         return nil
 210 }
 211
 212 func (rr *roundRobin) Start(target string, config BalancerConfig) error {
 213         rr.mu.Lock()
 214         defer rr.mu.Unlock()
 215         if rr.done {
 216                 return ErrClientConnClosing
 217         }
 218         if rr.r == nil {
 219                 // If there is no name resolver installed, it is not needed to
 220                 // do name resolution. In this case, target is added into rr.addrs
 221                 // as the only address available and rr.addrCh stays nil.
 222                 rr.addrs = append(rr.addrs, &addrInfo{addr: Address{Addr: target}})
 223                 return nil
 224         }
 225         w, err := rr.r.Resolve(target)
 226         if err != nil {
 227                 return err
 228         }
 229         rr.w = w
 230         rr.addrCh = make(chan []Address, 1)
 231         go func() {
 232                 for {
 233                         if err := rr.watchAddrUpdates(); err != nil {
 234                                 return
 235                         }
 236                 }
 237         }()
 238         return nil
 239 }
 240
 241 // Up sets the connected state of addr and sends notification if there are pending
 242 // Get() calls.
 243 func (rr *roundRobin) Up(addr Address) func(error) {
 244         rr.mu.Lock()
 245         defer rr.mu.Unlock()
 246         var cnt int
 247         for _, a := range rr.addrs {
 248                 if a.addr == addr {
 249                         if a.connected {
 250                                 return nil
 251                         }
 252                         a.connected = true
 253                 }
 254                 if a.connected {
 255                         cnt++
 256                 }
 257         }
 258         // addr is only one which is connected. Notify the Get() callers who are blocking.
 259         if cnt == 1 && rr.waitCh != nil {
 260                 close(rr.waitCh)
 261                 rr.waitCh = nil
 262         }
 263         return func(err error) {
 264                 rr.down(addr, err)
 265         }
 266 }
 267
 268 // down unsets the connected state of addr.
 269 func (rr *roundRobin) down(addr Address, err error) {
 270         rr.mu.Lock()
 271         defer rr.mu.Unlock()
 272         for _, a := range rr.addrs {
 273                 if addr == a.addr {
 274                         a.connected = false
 275                         break
 276                 }
 277         }
 278 }
 279
 280 // Get returns the next addr in the rotation.
 281 func (rr *roundRobin) Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error) {
 282         var ch chan struct{}
 283         rr.mu.Lock()
 284         if rr.done {
 285                 rr.mu.Unlock()
 286                 err = ErrClientConnClosing
 287                 return
 288         }
 289
 290         if len(rr.addrs) > 0 {
 291                 if rr.next >= len(rr.addrs) {
 292                         rr.next = 0
 293                 }
 294                 next := rr.next
 295                 for {
 296                         a := rr.addrs[next]
 297                         next = (next + 1) % len(rr.addrs)
 298                         if a.connected {
 299                                 addr = a.addr
 300                                 rr.next = next
 301                                 rr.mu.Unlock()
 302                                 return
 303                         }
 304                         if next == rr.next {
 305                                 // Has iterated all the possible address but none is connected.
 306                                 break
 307                         }
 308                 }
 309         }
 310         if !opts.BlockingWait {
 311                 if len(rr.addrs) == 0 {
 312                         rr.mu.Unlock()
 313                         err = Errorf(codes.Unavailable, "there is no address available")
 314                         return
 315                 }
 316                 // Returns the next addr on rr.addrs for failfast RPCs.
 317                 addr = rr.addrs[rr.next].addr
 318                 rr.next++
 319                 rr.mu.Unlock()
 320                 return
 321         }
 322         // Wait on rr.waitCh for non-failfast RPCs.
 323         if rr.waitCh == nil {
 324                 ch = make(chan struct{})
 325                 rr.waitCh = ch
 326         } else {
 327                 ch = rr.waitCh
 328         }
 329         rr.mu.Unlock()
 330         for {
 331                 select {
 332                 case <-ctx.Done():
 333                         err = ctx.Err()
 334                         return
 335                 case <-ch:
 336                         rr.mu.Lock()
 337                         if rr.done {
 338                                 rr.mu.Unlock()
 339                                 err = ErrClientConnClosing
 340                                 return
 341                         }
 342
 343                         if len(rr.addrs) > 0 {
 344                                 if rr.next >= len(rr.addrs) {
 345                                         rr.next = 0
 346                                 }
 347                                 next := rr.next
 348                                 for {
 349                                         a := rr.addrs[next]
 350                                         next = (next + 1) % len(rr.addrs)
 351                                         if a.connected {
 352                                                 addr = a.addr
 353                                                 rr.next = next
 354                                                 rr.mu.Unlock()
 355                                                 return
 356                                         }
 357                                         if next == rr.next {
 358                                                 // Has iterated all the possible address but none is connected.
 359                                                 break
 360                                         }
 361                                 }
 362                         }
 363                         // The newly added addr got removed by Down() again.
 364                         if rr.waitCh == nil {
 365                                 ch = make(chan struct{})
 366                                 rr.waitCh = ch
 367                         } else {
 368                                 ch = rr.waitCh
 369                         }
 370                         rr.mu.Unlock()
 371                 }
 372         }
 373 }
 374
 375 func (rr *roundRobin) Notify() <-chan []Address {
 376         return rr.addrCh
 377 }
 378
 379 func (rr *roundRobin) Close() error {
 380         rr.mu.Lock()
 381         defer rr.mu.Unlock()
 382         if rr.done {
 383                 return errBalancerClosed
 384         }
 385         rr.done = true
 386         if rr.w != nil {
 387                 rr.w.Close()
 388         }
 389         if rr.waitCh != nil {
 390                 close(rr.waitCh)
 391                 rr.waitCh = nil
 392         }
 393         if rr.addrCh != nil {
 394                 close(rr.addrCh)
 395         }
 396         return nil
 397 }
 398
 399 // pickFirst is used to test multi-addresses in one addrConn in which all addresses share the same addrConn.
 400 // It is a wrapper around roundRobin balancer. The logic of all methods works fine because balancer.Get()
 401 // returns the only address Up by resetTransport().
 402 type pickFirst struct {
 403         *roundRobin
 404 }
 405
 406 func pickFirstBalancerV1(r naming.Resolver) Balancer {
 407         return &pickFirst{&roundRobin{r: r}}
 408 }