fix(nacos): client.WithTag does not work (#83)

hertz-contrib · Jun 19, 2023 · fa188fd · fa188fd
1 parent 273e90d
commit fa188fd
Show file tree

Hide file tree

Showing 2 changed files with 203 additions and 59 deletions.
diff --git a/nacos/nacos_test.go b/nacos/nacos_test.go
@@ -23,7 +23,6 @@ import (
 	"github.com/cloudwego/hertz/pkg/app"
 	"github.com/cloudwego/hertz/pkg/app/client"
 	"github.com/cloudwego/hertz/pkg/app/client/discovery"
-	"github.com/cloudwego/hertz/pkg/app/client/loadbalance"
 	"github.com/cloudwego/hertz/pkg/app/middlewares/client/sd"
 	"github.com/cloudwego/hertz/pkg/app/server"
 	"github.com/cloudwego/hertz/pkg/app/server/registry"
@@ -361,56 +360,6 @@ func TestDefaultNacosRegistry(t *testing.T) {
 	assert.Nil(t, err)
 }
 
-// TestHertzAppWithNacosRegistry test a client call a hertz app with NacosRegistry
-func TestHertzAppWithNacosRegistry(t *testing.T) {
-	register := NewNacosRegistry(namingClient)
-	address := "127.0.0.1:4576"
-	srvName := "demo.hertz-contrib.testing"
-	var opts []config.Option
-	opts = append(opts, server.WithHostPorts(address))
-	opts = append(opts, server.WithRegistry(register, &registry.Info{
-		ServiceName: srvName,
-		Addr:        utils.NewNetAddr("tcp", address),
-		Weight:      10,
-		Tags:        nil,
-	}))
-	// run a hertz app,registry src info into NacosRegistry
-	srv := server.New(opts...)
-	srv.GET("/ping", func(c context.Context, ctx *app.RequestContext) {
-		ctx.String(200, "pong")
-	})
-	go srv.Spin()
-	// Because delayed registration, we need sleep more time.
-	time.Sleep(2 * time.Second)
-
-	// client call an url, with NacosResolver
-	newClient, _ := client.NewClient()
-	resolver := NewNacosResolver(namingClient)
-	newClient.Use(sd.Discovery(resolver, sd.WithLoadBalanceOptions(
-		loadbalance.NewWeightedBalancer(),
-		loadbalance.Options{
-			ExpireInterval:  3 * time.Second,
-			RefreshInterval: 1 * time.Second,
-		}),
-	))
-
-	status, body, err := newClient.Get(context.TODO(), nil, "http://demo.hertz-contrib.testing/ping",
-		config.WithSD(true))
-	assert.Nil(t, err)
-	assert.Equal(t, 200, status)
-	assert.Equal(t, "pong", string(body))
-
-	if err = srv.Shutdown(context.TODO()); err != nil {
-		t.Error(err)
-	}
-	time.Sleep(6 * time.Second)
-	status1, body1, err1 := newClient.Get(context.Background(), nil, "http://demo.hertz-contrib.testing/ping",
-		config.WithSD(true))
-	assert.NotNil(t, err1)
-	assert.Equal(t, 0, status1)
-	assert.Equal(t, "", string(body1))
-}
-
 // TestResolverDifferentGroup test NewResolver WithCluster option
 func TestResolverDifferentGroup(t *testing.T) {
 	var opts1 []config.Option
@@ -476,3 +425,145 @@ func TestResolverDifferentGroup(t *testing.T) {
 	assert.Equal(t, 200, status2)
 	assert.Equal(t, "pong2", string(body2))
 }
+
+func TestWithTag(t *testing.T) {
+	var opts1 []config.Option
+	var opts2 []config.Option
+
+	opts1 = append(opts1, server.WithRegistry(NewNacosRegistry(namingClient), &registry.Info{
+		ServiceName: "demo.hertz-contrib.test1",
+		Addr:        utils.NewNetAddr("tcp", "127.0.0.1:7512"),
+		Weight:      10,
+		Tags:        map[string]string{"key1": "value1"},
+	}))
+	opts1 = append(opts1, server.WithHostPorts("127.0.0.1:7512"))
+	srv1 := server.New(opts1...)
+	srv1.GET("/ping", func(c context.Context, ctx *app.RequestContext) {
+		ctx.String(200, "pong1")
+	})
+
+	opts2 = append(opts2, server.WithRegistry(NewNacosRegistry(namingClient), &registry.Info{
+		ServiceName: "demo.hertz-contrib.test1",
+		Addr:        utils.NewNetAddr("tcp", "127.0.0.1:7074"),
+		Weight:      10,
+		Tags:        map[string]string{"key2": "value2"},
+	}))
+	opts2 = append(opts2, server.WithHostPorts("127.0.0.1:7074"))
+	srv2 := server.New(opts2...)
+	srv2.GET("/ping", func(c context.Context, ctx *app.RequestContext) {
+		ctx.String(200, "pong2")
+	})
+
+	go srv1.Spin()
+	go srv2.Spin()
+
+	time.Sleep(2 * time.Second)
+
+	cli, _ := client.NewClient()
+	r := NewNacosResolver(namingClient)
+	cli.Use(sd.Discovery(r))
+
+	ctx, cancelFunc := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancelFunc()
+
+	status, body, err := cli.Get(ctx, nil,
+		"http://demo.hertz-contrib.test1/ping",
+		config.WithSD(true),
+		config.WithTag("key1", "value1"),
+	)
+	assert.Nil(t, err)
+	assert.Equal(t, 200, status)
+	assert.Equal(t, "pong1", string(body))
+}
+
+// TestCompareMaps tests the compareMaps function
+func TestCompareMaps(t *testing.T) {
+	// create some test cases with expected results
+	testCases := []struct {
+		m1, m2 map[string]string
+		want   bool
+	}{
+		{
+			m1:   map[string]string{"a": "1", "b": "2", "c": "3"},
+			m2:   map[string]string{"a": "1", "b": "2", "c": "3"},
+			want: true,
+		},
+		{
+			m1:   map[string]string{"a": "1", "b": "2", "c": "3"},
+			m2:   map[string]string{"a": "1", "b": "2", "d": "3"},
+			want: false,
+		},
+		{
+			m1:   map[string]string{"a": "1", "b": "2", "c": "3"},
+			m2:   map[string]string{"a": "1", "b": "2", "c": "4"},
+			want: false,
+		},
+		{
+			m1:   map[string]string{"a": "1", "b": "2"},
+			m2:   map[string]string{"a": "1", "b": "2", "c": "3"},
+			want: false,
+		},
+		{
+			m1:   nil,
+			m2:   nil,
+			want: true,
+		},
+		{
+			m1:   nil,
+			m2:   make(map[string]string),
+			want: true,
+		},
+	}
+	// iterate over the test cases and check if the function returns the expected result
+	for _, tc := range testCases {
+		got := compareMaps(tc.m1, tc.m2)
+		if got != tc.want {
+			t.Errorf("compareMaps(%v, %v) = %v, want %v", tc.m1, tc.m2, got, tc.want)
+		}
+	}
+}
+
+// TestHertzAppWithNacosRegistry test a client call a hertz app with NacosRegistry
+func TestHertzAppWithNacosRegistry(t *testing.T) {
+	register := NewNacosRegistry(namingClient)
+	address := "127.0.0.1:4576"
+	srvName := "d.h.t"
+	var opts []config.Option
+	opts = append(opts, server.WithHostPorts(address), server.WithExitWaitTime(2*time.Second))
+	opts = append(opts, server.WithRegistry(register, &registry.Info{
+		ServiceName: srvName,
+		Addr:        utils.NewNetAddr("tcp", address),
+		Weight:      10,
+		Tags:        nil,
+	}))
+	// run a hertz app,registry src info into NacosRegistry
+	srv := server.New(opts...)
+	srv.GET("/ping", func(c context.Context, ctx *app.RequestContext) {
+		ctx.String(200, "pong")
+	})
+	go srv.Spin()
+	// Because delayed registration, we need sleep more time.
+	time.Sleep(2 * time.Second)
+
+	// client call an url, with NacosResolver
+	newClient, _ := client.NewClient()
+	resolver := NewNacosResolver(namingClient)
+	newClient.Use(sd.Discovery(resolver))
+
+	status, body, err := newClient.Get(context.TODO(), nil, "http://d.h.t/ping",
+		config.WithSD(true))
+	assert.Nil(t, err)
+	assert.Equal(t, 200, status)
+	assert.Equal(t, "pong", string(body))
+
+	ctx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancelFunc()
+	srv.Shutdown(ctx) //nolint:errcheck // ignore error
+
+	time.Sleep(5 * time.Second)
+	status, body, err = newClient.Get(context.Background(), nil, "http://d.h.t/ping",
+		config.WithSD(true))
+	assert.NotNil(t, err)
+	assert.Equal(t, 0, status)
+	assert.Equal(t, "", string(body))
+}
diff --git a/nacos/resolver.go b/nacos/resolver.go
@@ -17,7 +17,9 @@ package nacos
 import (
 	"context"
 	"net"
+	"net/url"
 	"strconv"
+	"strings"
 
 	"github.com/cloudwego/hertz/pkg/app/client/discovery"
 	"github.com/hertz-contrib/registry/nacos/common"
@@ -57,31 +59,62 @@ func WithResolverGroup(group string) ResolverOption {
 }
 
 func (n *nacosResolver) Target(_ context.Context, target *discovery.TargetInfo) string {
-	return target.Host
+	var metadata strings.Builder
+
+	// Set serviceName and metadata to desc
+	tags := target.Tags
+	if len(tags) == 0 {
+		return target.Host
+	}
+
+	metadata.WriteString(target.Host)
+	metadata.WriteString("?")
+	values := url.Values{}
+	for k, v := range tags {
+		values.Add(k, v)
+	}
+	metadata.WriteString(values.Encode())
+	return metadata.String()
 }
 
 func (n *nacosResolver) Resolve(_ context.Context, desc string) (discovery.Result, error) {
+	var metadata map[string]string
+	serviceName := desc
+
+	// Get serviceName and metadata from desc
+	if strings.Contains(desc, "?") {
+		queries, _ := url.Parse(desc)
+		tags, _ := url.ParseQuery(queries.Query().Encode())
+
+		result := make(map[string]string)
+		for key, value := range tags {
+			result[key] = value[0]
+		}
+		metadata = result
+		serviceName = strings.Split(desc, "?")[0]
+	}
+
 	res, err := n.client.SelectInstances(vo.SelectInstancesParam{
-		ServiceName: desc,
+		ServiceName: serviceName,
 		HealthyOnly: true,
 		GroupName:   n.opts.group,
 		Clusters:    []string{n.opts.cluster},
 	})
 	if err != nil {
 		return discovery.Result{}, err
 	}
-
 	instances := make([]discovery.Instance, 0, len(res))
-	for _, in := range res {
-		if !in.Enable {
+	for _, ins := range res {
+		if !ins.Enable || !compareMaps(ins.Metadata, metadata) {
 			continue
 		}
-		formatPort := strconv.FormatUint(in.Port, 10)
+
+		formatPort := strconv.FormatUint(ins.Port, 10)
 		instances = append(instances,
 			discovery.NewInstance(
 				"tcp",
-				net.JoinHostPort(in.Ip, formatPort),
-				int(in.Weight), in.Metadata,
+				net.JoinHostPort(ins.Ip, formatPort),
+				int(ins.Weight), ins.Metadata,
 			),
 		)
 	}
@@ -116,3 +149,23 @@ func NewNacosResolver(cli naming_client.INamingClient, opts ...ResolverOption) d
 	}
 	return &nacosResolver{client: cli, opts: opt}
 }
+
+// compareMaps compares two maps regardless of nil or empty
+func compareMaps(m1, m2 map[string]string) bool {
+	// if both maps are nil, they are equal
+	if m1 == nil && m2 == nil {
+		return true
+	}
+	// if the lengths are different, the maps are not equal
+	if len(m1) != len(m2) {
+		return false
+	}
+	// iterate over the keys of m1 and check if they exist in m2 with the same value
+	for k, v := range m1 {
+		if v2, ok := m2[k]; !ok || v != v2 {
+			return false
+		}
+	}
+	// return true if no differences are found
+	return true
+}