Home | History | Annotate | Download | only in grpc 
      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 	"golang.org/x/net/context"
     23 	"google.golang.org/grpc/balancer"
     24 	"google.golang.org/grpc/connectivity"
     25 	"google.golang.org/grpc/grpclog"
     26 	"google.golang.org/grpc/resolver"
     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 	cc balancer.ClientConn
     48 	sc balancer.SubConn
     49 }
     50 
     51 func (b *pickfirstBalancer) HandleResolvedAddrs(addrs []resolver.Address, err error) {
     52 	if err != nil {
     53 		grpclog.Infof("pickfirstBalancer: HandleResolvedAddrs called with error %v", err)
     54 		return
     55 	}
     56 	if b.sc == nil {
     57 		b.sc, err = b.cc.NewSubConn(addrs, balancer.NewSubConnOptions{})
     58 		if err != nil {
     59 			grpclog.Errorf("pickfirstBalancer: failed to NewSubConn: %v", err)
     60 			return
     61 		}
     62 		b.cc.UpdateBalancerState(connectivity.Idle, &picker{sc: b.sc})
     63 		b.sc.Connect()
     64 	} else {
     65 		b.sc.UpdateAddresses(addrs)
     66 		b.sc.Connect()
     67 	}
     68 }
     69 
     70 func (b *pickfirstBalancer) HandleSubConnStateChange(sc balancer.SubConn, s connectivity.State) {
     71 	grpclog.Infof("pickfirstBalancer: HandleSubConnStateChange: %p, %v", sc, s)
     72 	if b.sc != sc {
     73 		grpclog.Infof("pickfirstBalancer: ignored state change because sc is not recognized")
     74 		return
     75 	}
     76 	if s == connectivity.Shutdown {
     77 		b.sc = nil
     78 		return
     79 	}
     80 
     81 	switch s {
     82 	case connectivity.Ready, connectivity.Idle:
     83 		b.cc.UpdateBalancerState(s, &picker{sc: sc})
     84 	case connectivity.Connecting:
     85 		b.cc.UpdateBalancerState(s, &picker{err: balancer.ErrNoSubConnAvailable})
     86 	case connectivity.TransientFailure:
     87 		b.cc.UpdateBalancerState(s, &picker{err: balancer.ErrTransientFailure})
     88 	}
     89 }
     90 
     91 func (b *pickfirstBalancer) Close() {
     92 }
     93 
     94 type picker struct {
     95 	err error
     96 	sc  balancer.SubConn
     97 }
     98 
     99 func (p *picker) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
    100 	if p.err != nil {
    101 		return nil, nil, p.err
    102 	}
    103 	return p.sc, nil, nil
    104 }
    105 
    106 func init() {
    107 	balancer.Register(newPickfirstBuilder())
    108 }
    109