From 608cf08ab4947ce23e60c67ff556dd52d71c88cf Mon Sep 17 00:00:00 2001
From: jizhuozhi <jizhuozhi.george@gmail.com>
Date: Thu, 2 Nov 2023 19:24:12 +0800
Subject: [PATCH 1/2] Cache picker only in taggingBalancer

---
 tagging.go      |  9 ---------
 tagging_test.go | 40 ++++++++++++++++++++++------------------
 2 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/tagging.go b/tagging.go
index adac867..8bacdae 100644
--- a/tagging.go
+++ b/tagging.go
@@ -65,9 +65,6 @@ 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))
 }
 
@@ -75,9 +72,6 @@ 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
index b890346..699e6b0 100644
--- a/tagging_test.go
+++ b/tagging_test.go
@@ -27,8 +27,6 @@ import (
 )
 
 type mockBalancer struct {
-	rebalanced bool
-	deleted    bool
 }
 
 type mockPicker struct {
@@ -39,14 +37,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 +88,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 +130,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 +170,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 +182,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 +210,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)

From a3eca062064db15a5f98d0418927a340f18de49b Mon Sep 17 00:00:00 2001
From: jizhuozhi <jizhuozhi.george@gmail.com>
Date: Thu, 2 Nov 2023 19:40:44 +0800
Subject: [PATCH 2/2] Cache picker only in taggingBalancer

---
 tagging_test.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tagging_test.go b/tagging_test.go
index 699e6b0..8599568 100644
--- a/tagging_test.go
+++ b/tagging_test.go
@@ -26,8 +26,7 @@ import (
 	"github.com/cloudwego/kitex/pkg/loadbalance"
 )
 
-type mockBalancer struct {
-}
+type mockBalancer struct{}
 
 type mockPicker struct {
 	result discovery.Result