grpc负载均衡RoundRobin源码解读

你的名字 2023-06-25 07:44 11阅读 0赞

grpc client端创建连接时可以用WithBalancer来指定负载均衡组件，这里研究下grpc自带的RoundRobin（轮询调度）的实现。源码在google.golang.org/grpc/balancer.go中。

roundRobin结构体定义如下：

type roundRobin struct {
    	r      naming.Resolver
    	w      naming.Watcher
    	addrs  []*addrInfo // all the addresses the client should potentially connect
    	mu     sync.Mutex
    	addrCh chan []Address // the channel to notify gRPC internals the list of addresses the client should connect to.
    	next   int            // index of the next address to return for Get()
    	waitCh chan struct{}  // the channel to block when there is no connected address available
    	done   bool           // The Balancer is closed.
    }

*  r是命名解析器，可以定义自己的命名解析器，如etcd命名解析器。如果r为nil，那么Dial中参数target将直接作为可请求地址添加到addrs中。
 *  w是命名解析器Resolve方法返回的watcher，该watcher可以监听命名解析器发来的地址信息变化，通知roundRobin对addrs中的地址进行动态的增删。
 *  addrs是从命名解析器获取地址信息数组，数组中每个地址不仅有地址信息，还有grpc与该地址是否已经创建了ready状态的连接。
 *  addrCh是地址数组的channel，该channel会在**每次**命名解析器发来地址信息变化后，将**所有**addrs通知到grpc内部的lbWatcher，lbWatcher是统一管理地址连接状态的协程，负责新地址的连接与被删除地址的关闭操作。
 *  next是roundRobin的Index，即轮询调度遍历到addrs数组中的哪个位置了。
 *  waitCh是当addrs中地址为空时，grpc调用Get()方法希望获取到一个到target的连接，如果设置了grpc的failfast为false，那么Get()方法会阻塞在此channel上，直到有ready的连接。

**roundRobin启动：**

func (rr *roundRobin) Start(target string, config BalancerConfig) error {
    	rr.mu.Lock()
    	defer rr.mu.Unlock()
    	if rr.done {
    		return ErrClientConnClosing
    	}
    	if rr.r == nil {
    		// 如果没有解析器，那么直接将target加入addrs地址数组
    		rr.addrs = append(rr.addrs, &addrInfo{addr: Address{Addr: target}})
    		return nil
    	}
    	// Resolve接口会返回一个watcher，watcher可以监听解析器的地址变化
    	w, err := rr.r.Resolve(target)
    	if err != nil {
    		return err
    	}
    	rr.w = w
    	// 创建一个channel，当watcher监听到地址变化时，通知grpc内部lbWatcher去连接该地址
    	rr.addrCh = make(chan []Address, 1)
    	// go 出去监听watcher，监听地址变化。
    	go func() {
    		for {
    			if err := rr.watchAddrUpdates(); err != nil {
    				return
    			}
    		}
    	}()
    	return nil
    }

**监听命名解析器的地址变化：**

func (rr *roundRobin) watchAddrUpdates() error {
    	// watcher的next方法会阻塞，直至有地址变化信息过来，updates即为变化信息
    	updates, err := rr.w.Next()
    	if err != nil {
    		return err
    	}
    	// 对于addrs地址数组的操作，显然是要加锁的，因为有多个goroutine在同时操作
    	rr.mu.Lock()
    	defer rr.mu.Unlock()
    	for _, update := range updates {
    		addr := Address{
    			Addr:     update.Addr,
    			Metadata: update.Metadata,
    		}
    		switch update.Op {
    		case naming.Add:
    		//对于新增类型的地址，注意这里不会重复添加。
    			var exist bool
    			for _, v := range rr.addrs {
    				if addr == v.addr {
    					exist = true
    					break
    				}
    			}
    			if exist {
    				continue
    			}
    			rr.addrs = append(rr.addrs, &addrInfo{addr: addr})
    		case naming.Delete:
    		//对于删除的地址，直接在addrs中删除就行了
    			for i, v := range rr.addrs {
    				if addr == v.addr {
    					copy(rr.addrs[i:], rr.addrs[i+1:])
    					rr.addrs = rr.addrs[:len(rr.addrs)-1]
    					break
    				}
    			}
    		default:
    			grpclog.Errorln("Unknown update.Op ", update.Op)
    		}
    	}
    	// 这里复制了整个addrs地址数组，然后丢到addrCh channel中通知grpc内部lbWatcher，
    	// lbWatcher会关闭删除的地址，连接新增的地址。
    	// 连接ready后会有专门的goroutine调用Up方法修改addrs中地址的状态。
    	open := make([]Address, len(rr.addrs))
    	for i, v := range rr.addrs {
    		open[i] = v.addr
    	}
    	if rr.done {
    		return ErrClientConnClosing
    	}
    	select {
    	case <-rr.addrCh:
    	default:
    	}
    	rr.addrCh <- open
    	return nil
    }

**Up方法：**  
up方法是grpc内部负载均衡watcher调用的，该watcher会读全局的连接状态改变队列，如果是ready状态的连接，会调用up方法来改变addrs地址数组中该地址的状态为**已连接**。

func (rr *roundRobin) Up(addr Address) func(error) {
    	rr.mu.Lock()
    	defer rr.mu.Unlock()
    	var cnt int
    	//将地址数组中的addr置为已连接状态，这样这个地址就可以被client使用了。
    	for _, a := range rr.addrs {
    		if a.addr == addr {
    			if a.connected {
    				return nil
    			}
    			a.connected = true
    		}
    		if a.connected {
    			cnt++
    		}
    	}
    	// 当有一个可用地址时，之前可能是0个，可能要很多client阻塞在获取连接地址上，这里通知所有的client有可用连接啦。
    	// 为什么只等于1时通知？因为可用地址数量>1时，client是不会阻塞的。
    	if cnt == 1 && rr.waitCh != nil {
    		close(rr.waitCh)
    		rr.waitCh = nil
    	}
    	//返回禁用该地址的方法
    	return func(err error) {
    		rr.down(addr, err)
    	}
    }

**down方法：**  
down方法就简单了, 直接找到addr置为不可用就行了。

//如果addr1已经被连接上了，但是resolver通知删除了，grpc内部如何处理关闭的逻辑？
    func (rr *roundRobin) down(addr Address, err error) {
    	rr.mu.Lock()
    	defer rr.mu.Unlock()
    	for _, a := range rr.addrs {
    		if addr == a.addr {
    			a.connected = false
    			break
    		}
    	}
    }

**Get()方法：**  
client 需要获取一个可用的地址，如果addrs为空，或者addrs不为空，但是地址都不可用（没连接），Get()方法会返回错误。但是如果设置了failfast = false，Get()方法会阻塞在waitCh channel上，直至Up方法给到通知，然后轮询调度可用的地址。

func (rr *roundRobin) Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error) {
    	var ch chan struct{}
    	rr.mu.Lock()
    	if rr.done {
    		rr.mu.Unlock()
    		err = ErrClientConnClosing
    		return
    	}
    
    	if len(rr.addrs) > 0 {
    		// addrs的长度可能变化，如果next值超出了，就置为0，从头开始调度。
    		if rr.next >= len(rr.addrs) {
    			rr.next = 0
    		}
    		next := rr.next
    		//遍历整个addrs数组，直到选出一个可用的地址
    		for {
    			a := rr.addrs[next]
    			// next值加一，当然是循环的，到len(addrs)后，变为0
    			next = (next + 1) % len(rr.addrs)
    			if a.connected {
    				addr = a.addr
    				rr.next = next
    				rr.mu.Unlock()
    				return
    			}
    			if next == rr.next {
    				// 遍历完一圈了，还没找到，走下面逻辑
    				break
    			}
    		}
    	}
    	if !opts.BlockingWait { //如果是非阻塞模式，如果没有可用地址，那么报错
    		if len(rr.addrs) == 0 {
    			rr.mu.Unlock()
    			err = status.Errorf(codes.Unavailable, "there is no address available")
    			return
    		}
    		// Returns the next addr on rr.addrs for failfast RPCs.
    		addr = rr.addrs[rr.next].addr
    		rr.next++
    		rr.mu.Unlock()
    		return
    	}
    	// Wait on rr.waitCh for non-failfast RPCs.
    	// 如果是阻塞模式，那么需要阻塞在waitCh上，直到Up方法给通知
    	if rr.waitCh == nil {
    		ch = make(chan struct{})
    		rr.waitCh = ch
    	} else {
    		ch = rr.waitCh
    	}
    	rr.mu.Unlock()
    	for {
    		select {
    		case <-ctx.Done():
    			err = ctx.Err()
    			return
    		case <-ch:
    			rr.mu.Lock()
    			if rr.done {
    				rr.mu.Unlock()
    				err = ErrClientConnClosing
    				return
    			}
    
    			if len(rr.addrs) > 0 {
    				if rr.next >= len(rr.addrs) {
    					rr.next = 0
    				}
    				next := rr.next
    				for {
    					a := rr.addrs[next]
    					next = (next + 1) % len(rr.addrs)
    					if a.connected {
    						addr = a.addr
    						rr.next = next
    						rr.mu.Unlock()
    						return
    					}
    					if next == rr.next {
    						// 遍历完一圈了，还没找到，可能刚Up的地址被down掉了，重新等待。
    						break
    					}
    				}
    			}
    			// The newly added addr got removed by Down() again.
    			if rr.waitCh == nil {
    				ch = make(chan struct{})
    				rr.waitCh = ch
    			} else {
    				ch = rr.waitCh
    			}
    			rr.mu.Unlock()
    		}
    	}
    }

lbWatcher：  
lbWatcher会监听地址变化信息，roundroubin每次有地址变化时，会将所有的地址通知给lbWatcher，lbWatcher本身维护了地址连接的map表，会找出新添加的地址和需要删除的地址，然后做连接、关闭操作，再调用roundRobin的Up/Down方法通知连接的状态。

func (bw *balancerWrapper) lbWatcher() {
    	notifyCh := bw.balancer.Notify()
    	if notifyCh == nil {
    		// 没有定义解析器，直接连接这个地址。
    		a := resolver.Address{
    			Addr: bw.targetAddr,
    			Type: resolver.Backend,
    		}
    		sc, err := bw.cc.NewSubConn([]resolver.Address{a}, balancer.NewSubConnOptions{})
    		if err != nil {
    			grpclog.Warningf("Error creating connection to %v. Err: %v", a, err)
    		} else {
    			bw.mu.Lock()
    			bw.conns[a] = sc
    			bw.connSt[sc] = &scState{
    				addr: Address{Addr: bw.targetAddr},
    				s:    connectivity.Idle,
    			}
    			bw.mu.Unlock()
    			sc.Connect()
    		}
    		return
    	}
    
    	for addrs := range notifyCh {
    			var newAddrs []resolver.Address
    			for _, a := range addrs {
    				newAddr := resolver.Address{
    					Addr:       a.Addr,
    					Type:       resolver.Backend, // All addresses from balancer are all backends.
    					ServerName: "",
    					Metadata:   a.Metadata,
    				}
    				newAddrs = append(newAddrs, newAddr)
    			}
    			var (
    				add []resolver.Address // Addresses need to setup connections.
    				del []balancer.SubConn // Connections need to tear down.
    			)
    			resAddrs := make(map[resolver.Address]Address)
    			for _, a := range addrs {
    				resAddrs[resolver.Address{
    					Addr:       a.Addr,
    					Type:       resolver.Backend, // All addresses from balancer are all backends.
    					ServerName: "",
    					Metadata:   a.Metadata,
    				}] = a
    			}
    			bw.mu.Lock()
    			// 新添加的地址，需要去新建连接
    			for a := range resAddrs {
    				if _, ok := bw.conns[a]; !ok {
    					add = append(add, a)
    				}
    			}
    			// 要被删除的地址，需要去关闭连接
    			for a, c := range bw.conns {
    				if _, ok := resAddrs[a]; !ok {
    					del = append(del, c)
    					delete(bw.conns, a)
    					// Keep the state of this sc in bw.connSt until its state becomes Shutdown.
    				}
    			}
    			bw.mu.Unlock()
    			for _, a := range add {
    				sc, err := bw.cc.NewSubConn([]resolver.Address{a}, balancer.NewSubConnOptions{})
    				if err != nil {
    					grpclog.Warningf("Error creating connection to %v. Err: %v", a, err)
    				} else {
    					bw.mu.Lock()
    					bw.conns[a] = sc
    					bw.connSt[sc] = &scState{
    						addr: resAddrs[a],
    						s:    connectivity.Idle,
    					}
    					bw.mu.Unlock()
    					sc.Connect()//  这一步真正做了连接的操作。
    				}
    			}
    			for _, c := range del {
    				bw.cc.RemoveSubConn(c)
    			}
    		}
    	}
    }