124
vendor/google.golang.org/grpc/grpclb_util.go
generated
vendored
124
vendor/google.golang.org/grpc/grpclb_util.go
generated
vendored
@@ -19,7 +19,12 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/grpc/balancer"
|
||||
"google.golang.org/grpc/connectivity"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
@@ -88,3 +93,122 @@ func (r *lbManualResolver) NewAddress(addrs []resolver.Address) {
|
||||
func (r *lbManualResolver) NewServiceConfig(sc string) {
|
||||
r.ccr.NewServiceConfig(sc)
|
||||
}
|
||||
|
||||
const subConnCacheTime = time.Second * 10
|
||||
|
||||
// lbCacheClientConn is a wrapper balancer.ClientConn with a SubConn cache.
|
||||
// SubConns will be kept in cache for subConnCacheTime before being removed.
|
||||
//
|
||||
// Its new and remove methods are updated to do cache first.
|
||||
type lbCacheClientConn struct {
|
||||
cc balancer.ClientConn
|
||||
timeout time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
// subConnCache only keeps subConns that are being deleted.
|
||||
subConnCache map[resolver.Address]*subConnCacheEntry
|
||||
subConnToAddr map[balancer.SubConn]resolver.Address
|
||||
}
|
||||
|
||||
type subConnCacheEntry struct {
|
||||
sc balancer.SubConn
|
||||
|
||||
cancel func()
|
||||
abortDeleting bool
|
||||
}
|
||||
|
||||
func newLBCacheClientConn(cc balancer.ClientConn) *lbCacheClientConn {
|
||||
return &lbCacheClientConn{
|
||||
cc: cc,
|
||||
timeout: subConnCacheTime,
|
||||
subConnCache: make(map[resolver.Address]*subConnCacheEntry),
|
||||
subConnToAddr: make(map[balancer.SubConn]resolver.Address),
|
||||
}
|
||||
}
|
||||
|
||||
func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
|
||||
if len(addrs) != 1 {
|
||||
return nil, fmt.Errorf("grpclb calling NewSubConn with addrs of length %v", len(addrs))
|
||||
}
|
||||
addrWithoutMD := addrs[0]
|
||||
addrWithoutMD.Metadata = nil
|
||||
|
||||
ccc.mu.Lock()
|
||||
defer ccc.mu.Unlock()
|
||||
if entry, ok := ccc.subConnCache[addrWithoutMD]; ok {
|
||||
// If entry is in subConnCache, the SubConn was being deleted.
|
||||
// cancel function will never be nil.
|
||||
entry.cancel()
|
||||
delete(ccc.subConnCache, addrWithoutMD)
|
||||
return entry.sc, nil
|
||||
}
|
||||
|
||||
scNew, err := ccc.cc.NewSubConn(addrs, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ccc.subConnToAddr[scNew] = addrWithoutMD
|
||||
return scNew, nil
|
||||
}
|
||||
|
||||
func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
|
||||
ccc.mu.Lock()
|
||||
defer ccc.mu.Unlock()
|
||||
addr, ok := ccc.subConnToAddr[sc]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if entry, ok := ccc.subConnCache[addr]; ok {
|
||||
if entry.sc != sc {
|
||||
// This could happen if NewSubConn was called multiple times for the
|
||||
// same address, and those SubConns are all removed. We remove sc
|
||||
// immediately here.
|
||||
delete(ccc.subConnToAddr, sc)
|
||||
ccc.cc.RemoveSubConn(sc)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
entry := &subConnCacheEntry{
|
||||
sc: sc,
|
||||
}
|
||||
ccc.subConnCache[addr] = entry
|
||||
|
||||
timer := time.AfterFunc(ccc.timeout, func() {
|
||||
ccc.mu.Lock()
|
||||
if entry.abortDeleting {
|
||||
return
|
||||
}
|
||||
ccc.cc.RemoveSubConn(sc)
|
||||
delete(ccc.subConnToAddr, sc)
|
||||
delete(ccc.subConnCache, addr)
|
||||
ccc.mu.Unlock()
|
||||
})
|
||||
entry.cancel = func() {
|
||||
if !timer.Stop() {
|
||||
// If stop was not successful, the timer has fired (this can only
|
||||
// happen in a race). But the deleting function is blocked on ccc.mu
|
||||
// because the mutex was held by the caller of this function.
|
||||
//
|
||||
// Set abortDeleting to true to abort the deleting function. When
|
||||
// the lock is released, the deleting function will acquire the
|
||||
// lock, check the value of abortDeleting and return.
|
||||
entry.abortDeleting = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ccc *lbCacheClientConn) UpdateBalancerState(s connectivity.State, p balancer.Picker) {
|
||||
ccc.cc.UpdateBalancerState(s, p)
|
||||
}
|
||||
|
||||
func (ccc *lbCacheClientConn) close() {
|
||||
ccc.mu.Lock()
|
||||
// Only cancel all existing timers. There's no need to remove SubConns.
|
||||
for _, entry := range ccc.subConnCache {
|
||||
entry.cancel()
|
||||
}
|
||||
ccc.mu.Unlock()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user