Cache picker only in taggingBalancer (#3)

kitex-contrib · Nov 2, 2023 · 19f1bc5 · 19f1bc5
1 parent 1b942bb
commit 19f1bc5
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 29 deletions.
diff --git a/tagging.go b/tagging.go
@@ -65,19 +65,13 @@ func (b *taggingBalancer) Rebalance(change discovery.Change) {
 	if !change.Result.Cacheable {
 		return
 	}
-	if next, ok := b.next.(loadbalance.Rebalancer); ok {
-		next.Delete(change)
-	}
 	b.pickerCache.Store(change.Result.CacheKey, b.createPicker(change.Result))
 }
 
 func (b *taggingBalancer) Delete(change discovery.Change) {
 	if !change.Result.Cacheable {
 		return
 	}
-	if next, ok := b.next.(loadbalance.Rebalancer); ok {
-		next.Delete(change)
-	}
 	b.pickerCache.Delete(change.Result.CacheKey)
 }
 
@@ -93,10 +87,7 @@ func (b *taggingBalancer) createPicker(e discovery.Result) loadbalance.Picker {
 
 	pickers := make(map[string]loadbalance.Picker, len(instances))
 	for t, instances := range instances {
-		// a projection of raw discovery.Result has same cache option
 		p := b.next.GetPicker(discovery.Result{
-			Cacheable: e.Cacheable,
-			CacheKey:  e.CacheKey,
 			Instances: instances,
 		})
 		pickers[t] = p

diff --git a/tagging_test.go b/tagging_test.go
@@ -26,10 +26,7 @@ import (
 	"github.com/cloudwego/kitex/pkg/loadbalance"
 )
 
-type mockBalancer struct {
-	rebalanced bool
-	deleted    bool
-}
+type mockBalancer struct{}
 
 type mockPicker struct {
 	result discovery.Result
@@ -39,14 +36,6 @@ func (m *mockBalancer) GetPicker(result discovery.Result) loadbalance.Picker {
 	return &mockPicker{result: result}
 }
 
-func (m *mockBalancer) Rebalance(change discovery.Change) {
-	m.rebalanced = true
-}
-
-func (m *mockBalancer) Delete(change discovery.Change) {
-	m.deleted = true
-}
-
 func (m *mockPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
 	return m.result.Instances[0]
 }
@@ -98,6 +87,26 @@ func TestTaggingBalancer_GetPicker(t *testing.T) {
 				},
 			},
 		},
+		{
+			cacheable: true,
+			cacheKey:  "multi instances",
+			instances: []discovery.Instance{
+				discovery.NewInstance("tcp", "addr1", 10, map[string]string{"foo": "bar1"}),
+				discovery.NewInstance("tcp", "addr2", 20, map[string]string{"foo": "bar2"}),
+				discovery.NewInstance("tcp", "addr3", 30, map[string]string{"foo": "bar3"}),
+				discovery.NewInstance("tcp", "addr4", 30, map[string]string{"foo": ""}),
+				discovery.NewInstance("tcp", "addr5", 30, nil),
+			},
+			tagInstances: map[string][]discovery.Instance{
+				"bar1": {discovery.NewInstance("tcp", "addr1", 10, map[string]string{"foo": "bar1"})},
+				"bar2": {discovery.NewInstance("tcp", "addr2", 20, map[string]string{"foo": "bar2"})},
+				"bar3": {discovery.NewInstance("tcp", "addr3", 30, map[string]string{"foo": "bar3"})},
+				"": {
+					discovery.NewInstance("tcp", "addr4", 30, map[string]string{"foo": ""}),
+					discovery.NewInstance("tcp", "addr5", 30, nil),
+				},
+			},
+		},
 	}
 
 	lb := New("foo", func(ctx context.Context, request interface{}) string {
@@ -120,8 +129,8 @@ func TestTaggingBalancer_GetPicker(t *testing.T) {
 			assert.IsType(t, &mockPicker{}, p)
 
 			pp := p.(*mockPicker)
-			assert.Equal(t, tt.cacheable, pp.result.Cacheable)
-			assert.Equal(t, tt.cacheKey, pp.result.CacheKey)
+			assert.Zero(t, pp.result.Cacheable)
+			assert.Zero(t, pp.result.CacheKey)
 			assert.EqualValues(t, v, pp.result.Instances)
 		}
 	}
@@ -160,7 +169,6 @@ func TestTaggingBalancer_Rebalance(t *testing.T) {
 			Instances: []discovery.Instance{discovery.NewInstance("tcp", "addr2", 20, map[string]string{"foo": "bar"})},
 		},
 	})
-	assert.True(t, lb.(*taggingBalancer).next.(*mockBalancer).deleted)
 
 	p2 := lb.GetPicker(discovery.Result{
 		Cacheable: true,
@@ -173,13 +181,9 @@ func TestTaggingBalancer_Rebalance(t *testing.T) {
 	mp1 := p1.(*taggingPicker).tagPickers["bar"].(*mockPicker)
 	mp2 := p2.(*taggingPicker).tagPickers["bar"].(*mockPicker)
 	assert.Equal(t, mp1.result, discovery.Result{
-		Cacheable: true,
-		CacheKey:  "rebalance",
 		Instances: []discovery.Instance{discovery.NewInstance("tcp", "addr1", 10, map[string]string{"foo": "bar"})},
 	})
 	assert.Equal(t, mp2.result, discovery.Result{
-		Cacheable: true,
-		CacheKey:  "rebalance",
 		Instances: []discovery.Instance{discovery.NewInstance("tcp", "addr2", 20, map[string]string{"foo": "bar"})},
 	})
 }
@@ -205,7 +209,6 @@ func TestTaggingBalancer_Delete(t *testing.T) {
 			CacheKey:  "delete",
 		},
 	})
-	assert.True(t, lb.(*taggingBalancer).next.(*mockBalancer).deleted)
 
 	pp, ok = lb.(*taggingBalancer).pickerCache.Load("delete")
 	assert.False(t, ok)