fix(lbcache): check the existence before new Balancer to prevent leakage (cloudwego#1825)

Author: ppzqh

pkg/loadbalance/lbcache/cache.go

@@ -163,6 +163,10 @@ func (b *BalancerFactory) Get(ctx context.Context, target rpcinfo.EndpointInfo)
 		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

pkg/loadbalance/lbcache/cache_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
@@ -126,3 +127,49 @@ func (m *mockRebalancer) Delete(ch discovery.Change) {
 		m.deleteFunc(ch)
 	}
 }
+
+type mockResolver struct{}
+
+func (m *mockResolver) Target(ctx context.Context, target rpcinfo.EndpointInfo) (description string) {
+	return "target"
+}
+
+func (m *mockResolver) Resolve(ctx context.Context, desc string) (discovery.Result, error) {
+	return discovery.Result{}, nil
+}
+
+func (m *mockResolver) Diff(cacheKey string, prev, next discovery.Result) (discovery.Change, bool) {
+	return discovery.Change{}, false
+}
+
+func (m *mockResolver) Name() string {
+	return "name"
+}
+
+var _ discovery.Resolver = &mockResolver{}
+
+func TestConcurrentGet(t *testing.T) {
+	cacheOpts := Options{Cacheable: false, RefreshInterval: time.Second, ExpireInterval: 5 * time.Second}
+	bf := newBalancerFactory(&mockResolver{}, loadbalance.NewWeightedBalancer(), cacheOpts)
+	m := sync.Map{}
+	wg := sync.WaitGroup{}
+
+	// concurrent get
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			b, _ := bf.Get(context.Background(), nil)
+			m.Store(b, b)
+		}()
+	}
+	wg.Wait()
+
+	// check if length == 1, (target -> balancer, 1:1)
+	cnt := 0
+	m.Range(func(key, value any) bool {
+		cnt++
+		return true
+	})
+	test.Assert(t, cnt == 1, cnt)
+}