feat: wrap singleflight.group to provide CheckAndDo

Author: ppzqh

pkg/loadbalance/consist.go

@@ -23,8 +23,6 @@ import (
 	"time"
 
 	"github.com/bytedance/gopkg/util/xxhash3"
-	"golang.org/x/sync/singleflight"
-
 	"github.com/cloudwego/kitex/pkg/discovery"
 	"github.com/cloudwego/kitex/pkg/utils"
 )
@@ -226,7 +224,7 @@ func buildConsistResult(info *consistInfo, key uint64) *consistResult {
 type consistBalancer struct {
 	cachedConsistInfo sync.Map
 	opt               ConsistentHashOption
-	sfg               singleflight.Group
+	sfg               utils.SingleFlightGroup
 }
 
 // NewConsistBalancer creates a new consist balancer with the given option.
@@ -247,13 +245,17 @@ func NewConsistBalancer(opt ConsistentHashOption) Loadbalancer {
 func (cb *consistBalancer) GetPicker(e discovery.Result) Picker {
 	var ci *consistInfo
 	if e.Cacheable {
-		cii, ok := cb.cachedConsistInfo.Load(e.CacheKey)
-		if !ok {
-			cii, _, _ = cb.sfg.Do(e.CacheKey, func() (interface{}, error) {
-				return cb.newConsistInfo(e), nil
-			})
-			cb.cachedConsistInfo.Store(e.CacheKey, cii)
-		}
+		cii, _, _ := cb.sfg.CheckAndDo(
+			e.CacheKey,
+			func() (any, bool) {
+				return cb.cachedConsistInfo.Load(e.CacheKey)
+			},
+			func() (interface{}, error) {
+				res := cb.newConsistInfo(e)
+				cb.cachedConsistInfo.Store(e.CacheKey, res)
+				return res, nil
+			},
+		)
 		ci = cii.(*consistInfo)
 	} else {
 		ci = cb.newConsistInfo(e)

pkg/loadbalance/lbcache/cache.go

@@ -24,8 +24,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"golang.org/x/sync/singleflight"
-
 	"github.com/cloudwego/kitex/pkg/diagnosis"
 	"github.com/cloudwego/kitex/pkg/discovery"
 	"github.com/cloudwego/kitex/pkg/klog"
@@ -41,7 +39,7 @@ const (
 
 var (
 	balancerFactories    sync.Map // key: resolver name + loadbalance name
-	balancerFactoriesSfg singleflight.Group
+	balancerFactoriesSfg utils.SingleFlightGroup
 )
 
 // Options for create builder
@@ -88,7 +86,7 @@ type BalancerFactory struct {
 	resolver   discovery.Resolver
 	balancer   loadbalance.Loadbalancer
 	rebalancer loadbalance.Rebalancer
-	sfg        singleflight.Group
+	sfg        utils.SingleFlightGroup
 }
 
 func cacheKey(resolver, balancer string, opts Options) string {
@@ -120,15 +118,15 @@ func NewBalancerFactory(resolver discovery.Resolver, balancer loadbalance.Loadba
 		return newBalancerFactory(resolver, balancer, opts)
 	}
 	uniqueKey := cacheKey(resolver.Name(), balancer.Name(), opts)
-	val, ok := balancerFactories.Load(uniqueKey)
-	if ok {
-		return val.(*BalancerFactory)
-	}
-	val, _, _ = balancerFactoriesSfg.Do(uniqueKey, func() (interface{}, error) {
-		b := newBalancerFactory(resolver, balancer, opts)
-		balancerFactories.Store(uniqueKey, b)
-		return b, nil
-	})
+	val, _, _ := balancerFactoriesSfg.CheckAndDo(uniqueKey,
+		func() (any, bool) {
+			return balancerFactories.Load(uniqueKey)
+		},
+		func() (interface{}, error) {
+			b := newBalancerFactory(resolver, balancer, opts)
+			balancerFactories.Store(uniqueKey, b)
+			return b, nil
+		})
 	return val.(*BalancerFactory)
 }
 
@@ -158,29 +156,25 @@ func renameResultCacheKey(res *discovery.Result, resolverName string) {
 // Get create a new balancer if not exists
 func (b *BalancerFactory) Get(ctx context.Context, target rpcinfo.EndpointInfo) (*Balancer, error) {
 	desc := b.resolver.Target(ctx, target)
-	val, ok := b.cache.Load(desc)
-	if ok {
-		return val.(*Balancer), nil
-	}
-	val, err, _ := b.sfg.Do(desc, func() (interface{}, error) {
-		if v, ok := b.cache.Load(desc); ok {
-			// cache may be set already
-			return v.(*Balancer), nil
-		}
-		res, err := b.resolver.Resolve(ctx, desc)
-		if err != nil {
-			return nil, err
-		}
-		renameResultCacheKey(&res, b.resolver.Name())
-		bl := &Balancer{
-			b:      b,
-			target: desc,
-		}
-		bl.res.Store(res)
-		bl.sharedTicker = getSharedTicker(bl, b.opts.RefreshInterval)
-		b.cache.Store(desc, bl)
-		return bl, nil
-	})
+	val, err, _ := b.sfg.CheckAndDo(desc,
+		func() (any, bool) {
+			return b.cache.Load(desc)
+		},
+		func() (interface{}, error) {
+			res, err := b.resolver.Resolve(ctx, desc)
+			if err != nil {
+				return nil, err
+			}
+			renameResultCacheKey(&res, b.resolver.Name())
+			bl := &Balancer{
+				b:      b,
+				target: desc,
+			}
+			bl.res.Store(res)
+			bl.sharedTicker = getSharedTicker(bl, b.opts.RefreshInterval)
+			b.cache.Store(desc, bl)
+			return bl, nil
+		})
 	if err != nil {
 		return nil, err
 	}

pkg/loadbalance/lbcache/shared_ticker.go

@@ -20,32 +20,27 @@ import (
 	"sync"
 	"time"
 
-	"golang.org/x/sync/singleflight"
-
 	"github.com/cloudwego/kitex/pkg/utils"
 )
 
 var (
 	// insert, not delete
 	sharedTickers    sync.Map
-	sharedTickersSfg singleflight.Group
+	sharedTickersSfg utils.SingleFlightGroup
 )
 
 func getSharedTicker(b *Balancer, refreshInterval time.Duration) *utils.SharedTicker {
-	sti, ok := sharedTickers.Load(refreshInterval)
-	if ok {
-		st := sti.(*utils.SharedTicker)
-		st.Add(b)
-		return st
-	}
-	v, _, _ := sharedTickersSfg.Do(refreshInterval.String(), func() (interface{}, error) {
-		st := utils.NewSharedTicker(refreshInterval)
-		sharedTickers.Store(refreshInterval, st)
-		return st, nil
-	})
+	v, _, _ := sharedTickersSfg.CheckAndDo(
+		refreshInterval.String(),
+		func() (any, bool) {
+			return sharedTickers.Load(refreshInterval)
+		},
+		func() (interface{}, error) {
+			st := utils.NewSharedTicker(refreshInterval)
+			sharedTickers.Store(refreshInterval, st)
+			return st, nil
+		})
 	st := v.(*utils.SharedTicker)
-	// Add without singleflight,
-	// because we need all refreshers those call this function to add themselves to SharedTicker
 	st.Add(b)
 	return st
 }

pkg/loadbalance/weighted_balancer.go

@@ -19,7 +19,7 @@ package loadbalance
 import (
 	"sync"
 
-	"golang.org/x/sync/singleflight"
+	"github.com/cloudwego/kitex/pkg/utils"
 
 	"github.com/cloudwego/kitex/pkg/discovery"
 	"github.com/cloudwego/kitex/pkg/klog"
@@ -35,7 +35,7 @@ const (
 type weightedBalancer struct {
 	kind        int
 	pickerCache sync.Map
-	sfg         singleflight.Group
+	sfg         utils.SingleFlightGroup
 }
 
 // NewWeightedBalancer creates a loadbalancer using weighted-round-robin algorithm.
@@ -73,15 +73,15 @@ func (wb *weightedBalancer) GetPicker(e discovery.Result) Picker {
 		picker := wb.createPicker(e)
 		return picker
 	}
-
-	picker, ok := wb.pickerCache.Load(e.CacheKey)
-	if !ok {
-		picker, _, _ = wb.sfg.Do(e.CacheKey, func() (interface{}, error) {
+	picker, _, _ := wb.sfg.CheckAndDo(e.CacheKey,
+		func() (any, bool) {
+			return wb.pickerCache.Load(e.CacheKey)
+		},
+		func() (interface{}, error) {
 			p := wb.createPicker(e)
 			wb.pickerCache.Store(e.CacheKey, p)
 			return p, nil
 		})
-	}
 	return picker.(Picker)
 }
 

pkg/remote/trans/nphttp2/conn_pool.go

@@ -26,7 +26,7 @@ import (
 	"sync/atomic"
 	"time"
 
-	"golang.org/x/sync/singleflight"
+	"github.com/cloudwego/kitex/pkg/utils"
 
 	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/remote"
@@ -57,7 +57,7 @@ func NewConnPool(remoteService string, size uint32, connOpts grpc.ConnectOptions
 // MuxPool manages a pool of long connections.
 type connPool struct {
 	size          uint32
-	sfg           singleflight.Group
+	sfg           utils.SingleFlightGroup
 	conns         sync.Map // key: address, value: *transports
 	remoteService string   // remote service name
 	connOpts      grpc.ConnectOptions
@@ -141,41 +141,39 @@ func (p *connPool) Get(ctx context.Context, network, address string, opt remote.
 
 	var (
 		trans *transports
-		conn  *clientConn
 		err   error
 	)
 
-	v, ok := p.conns.Load(address)
-	if ok {
-		trans = v.(*transports)
-		if tr := trans.get(); tr != nil {
-			if tr.(grpc.IsActive).IsActive() {
-				// Actually new a stream, reuse the connection (grpc.ClientTransport)
-				conn, err = newClientConn(ctx, tr, address)
-				if err == nil {
-					return conn, nil
+	tr, err, _ := p.sfg.CheckAndDo(address,
+		func() (any, bool) {
+			v, ok := p.conns.Load(address)
+			if ok {
+				trans = v.(*transports)
+				if tr := trans.get(); tr != nil {
+					if tr.(grpc.IsActive).IsActive() {
+						return tr, true
+					}
 				}
-				klog.CtxDebugf(ctx, "KITEX: New grpc stream failed, network=%s, address=%s, error=%s", network, address, err.Error())
 			}
-		}
-	}
-	tr, err, _ := p.sfg.Do(address, func() (i interface{}, e error) {
-		// Notice: newTransport means new a connection, the timeout of connection cannot be set,
-		// so using context.Background() but not the ctx passed in as the parameter.
-		tr, err := p.newTransport(context.Background(), opt.Dialer, network, address, opt.ConnectTimeout, p.connOpts)
-		if err != nil {
-			return nil, err
-		}
-		if trans == nil {
-			trans = &transports{
-				size:          p.size,
-				cliTransports: make([]grpc.ClientTransport, p.size),
+			return nil, false
+		},
+		func() (i interface{}, e error) {
+			// Notice: newTransport means new a connection, the timeout of connection cannot be set,
+			// so using context.Background() but not the ctx passed in as the parameter.
+			tr, err := p.newTransport(context.Background(), opt.Dialer, network, address, opt.ConnectTimeout, p.connOpts)
+			if err != nil {
+				return nil, err
 			}
-		}
-		trans.put(tr) // the tr (connection) maybe not in the pool, but can be recycled by keepalive.
-		p.conns.Store(address, trans)
-		return tr, nil
-	})
+			if trans == nil {
+				trans = &transports{
+					size:          p.size,
+					cliTransports: make([]grpc.ClientTransport, p.size),
+				}
+			}
+			trans.put(tr) // the tr (connection) maybe not in the pool, but can be recycled by keepalive.
+			p.conns.Store(address, trans)
+			return tr, nil
+		})
 	if err != nil {
 		klog.CtxErrorf(ctx, "KITEX: New grpc client connection failed, network=%s, address=%s, error=%s", network, address, err.Error())
 		return nil, err

pkg/utils/singleFlight.go

@@ -0,0 +1,24 @@
+package utils
+
+import "golang.org/x/sync/singleflight"
+
+// SingleFlightGroup is a wrapper around singleflight.Group to provide a CheckAndDo functionality.
+// It is used to ensure that `fn` is executed only once for a given key.
+type SingleFlightGroup struct {
+	singleflight.Group
+}
+
+// CheckAndDo implements a double-checked pattern to ensure that `fn` is executed only once for a given key.
+// It performs an initial check with `checkFunc` before calling `singleflight.Do`. If the value is found, return.
+// Otherwise, call `Do` and execute `checkFunc` again before executing `fn`.
+func (g *SingleFlightGroup) CheckAndDo(key string, checkFunc func() (any, bool), fn func() (any, error)) (v any, err error, shared bool) {
+	if value, loaded := checkFunc(); loaded {
+		return value, nil, true
+	}
+	return g.Do(key, func() (any, error) {
+		if value, loaded := checkFunc(); loaded {
+			return value, nil
+		}
+		return fn()
+	})
+}