Repository: go-kratos/kratos
Pull request #3509 (2024-12-27 08:19:06)

refactor(contrib/registry/etcd): use service ID as Registry.ctxMap key instead of service pointer
The code modification is based on two points:  
- It is not recommended to use pointer as key in Go map.
- Changing the service ID to be the key in ctxMap offers better performance. Detailed performance tests will be provided below.

test_kratos_registry.go
```go
package test_kratos_registry

import (
	"context"
	"testing"
	"time"
)

type ServiceInstance struct {
	ID        string            `json:"id"`
	Name      string            `json:"name"`
	Version   string            `json:"version"`
	Metadata  map[string]string `json:"metadata"`
	Endpoints []string          `json:"endpoints"`
}

type options struct {
	ctx       context.Context
	namespace string
	ttl       time.Duration
	maxRetry  int
}

type Registry struct {
	opts   *options
	ctxMap map[*ServiceInstance]context.CancelFunc
}

func (r *Registry) Register(ctx context.Context, service *ServiceInstance) {
	_, cancel := context.WithCancel(r.opts.ctx)
	r.ctxMap[service] = cancel
}

func (r *Registry) Deregister(ctx context.Context, service *ServiceInstance) {
	if cancel, ok := r.ctxMap[service]; ok {
		cancel()
		delete(r.ctxMap, service)
	}
}

type serviceCancel struct {
	svc    *ServiceInstance
	cancel context.CancelFunc
}

type RegistryNew struct {
	opts   *options
	ctxMap map[string]*serviceCancel
}

func (r *RegistryNew) Register(ctx context.Context, service *ServiceInstance) {
	_, cancel := context.WithCancel(r.opts.ctx)
	r.ctxMap[service.ID] = &serviceCancel{
		svc:    service,
		cancel: cancel,
	}
}

func (r *RegistryNew) Deregister(ctx context.Context, service *ServiceInstance) {
	if svcCancel, ok := r.ctxMap[service.ID]; ok {
		svcCancel.cancel()
		delete(r.ctxMap, service.ID)
	}
}

func newServiceInstances() []*ServiceInstance {
	return []*ServiceInstance{
		{
			ID:        "1",
			Name:      "service1",
			Version:   "v1",
			Metadata:  map[string]string{"key1": "value1"},
			Endpoints: []string{"http://addr1:port1"},
		},
		{
			ID:        "2",
			Name:      "service2",
			Version:   "v2",
			Metadata:  map[string]string{"key2": "value2"},
			Endpoints: []string{"http://addr2:port2"},
		},
		{
			ID:        "3",
			Name:      "service3",
			Version:   "v3",
			Metadata:  map[string]string{"key3": "value3"},
			Endpoints: []string{"http://addr3:port3"},
		},
		{
			ID:        "4",
			Name:      "service4",
			Version:   "v4",
			Metadata:  map[string]string{"key4": "value4"},
			Endpoints: []string{"http://addr4:port4"},
		},
	}
}

func BenchmarkRegisterRegister(b *testing.B) {
	r := &Registry{
		opts: &options{
			ctx: context.Background(),
		},
		ctxMap: make(map[*ServiceInstance]context.CancelFunc),
	}
	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		for _, svc := range newServiceInstances() {
			r.Register(context.Background(), svc)
			//break
		}
	}
}

func BenchmarkRegisterNewRegister(b *testing.B) {
	r := &RegistryNew{
		opts: &options{
			ctx: context.Background(),
		},
		ctxMap: make(map[string]*serviceCancel),
	}
	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		for _, svc := range newServiceInstances() {
			r.Register(context.Background(), svc)
			//break
		}
	}
}

func BenchmarkRegisterDeregister(b *testing.B) {
	r := &Registry{
		opts: &options{
			ctx: context.Background(),
		},
		ctxMap: make(map[*ServiceInstance]context.CancelFunc),
	}
	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		for _, svc := range newServiceInstances() {
			r.Deregister(context.Background(), svc)
			//break
		}
	}
}

func BenchmarkRegisterNewDeregister(b *testing.B) {
	r := &RegistryNew{
		opts: &options{
			ctx: context.Background(),
		},
		ctxMap: make(map[string]*serviceCancel),
	}
	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		for _, svc := range newServiceInstances() {
			r.Deregister(context.Background(), svc)
			//break
		}
	}
}
```

benchmark output:
```text
goos: darwin
goarch: arm64
pkg: test_code/test_kratos_registry
cpu: Apple M3 Pro
BenchmarkRegisterRegister
BenchmarkRegisterRegister-11         	 1000000	      1814 ns/op	    2454 B/op	      24 allocs/op
BenchmarkRegisterNewRegister
BenchmarkRegisterNewRegister-11      	 1783624	       764.8 ns/op	    2176 B/op	      28 allocs/op
BenchmarkRegisterDeregister
BenchmarkRegisterDeregister-11       	16407162	        71.50 ns/op	       0 B/op	       0 allocs/op
BenchmarkRegisterNewDeregister
BenchmarkRegisterNewDeregister-11    	17581944	        69.60 ns/op	       0 B/op	       0 allocs/op
PASS

Process finished with the exit code 0
```


Patch:
diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index 9e87d307..977e4ea3 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -44,7 +44,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up Go
-        uses: actions/setup-go@v5.1.0
+        uses: actions/setup-go@v5.4.0
         with:
           go-version: ${{ matrix.go }}
 
diff --git a/contrib/polaris/config.go b/contrib/polaris/config.go
index d4b71bd3..f6dbd38f 100644
--- a/contrib/polaris/config.go
+++ b/contrib/polaris/config.go
@@ -40,20 +40,20 @@ type source struct {
 
 // Load return the config values
 func (s *source) Load() ([]*config.KeyValue, error) {
-	kv := make([]*config.KeyValue, 0)
+	kvs := make([]*config.KeyValue, 0, len(s.options.files))
 	for _, file := range s.options.files {
 		configFile, err := s.client.GetConfigFile(s.options.namespace, file.Group, file.Name)
 		if err != nil {
 			return nil, err
 		}
 		s.options.configFile = append(s.options.configFile, configFile)
-		kv = append(kv, &config.KeyValue{
+		kvs = append(kvs, &config.KeyValue{
 			Key:    file.Name,
 			Value:  []byte(configFile.GetContent()),
 			Format: strings.TrimPrefix(filepath.Ext(file.Name), "."),
 		})
 	}
-	return kv, nil
+	return kvs, nil
 }
 
 // Watch return the watcher
diff --git a/contrib/registry/etcd/registry.go b/contrib/registry/etcd/registry.go
index 1a931153..61402ac3 100644
--- a/contrib/registry/etcd/registry.go
+++ b/contrib/registry/etcd/registry.go
@@ -55,7 +55,12 @@ type Registry struct {
 		ctxMap is used to store the context cancel function of each service instance.
 		When the service instance is deregistered, the corresponding context cancel function is called to stop the heartbeat.
 	*/
-	ctxMap map[*registry.ServiceInstance]context.CancelFunc
+	ctxMap map[string]*serviceCancel
+}
+
+type serviceCancel struct {
+	service *registry.ServiceInstance
+	cancel  context.CancelFunc
 }
 
 // New creates etcd registry
@@ -73,7 +78,7 @@ func New(client *clientv3.Client, opts ...Option) (r *Registry) {
 		opts:   op,
 		client: client,
 		kv:     clientv3.NewKV(client),
-		ctxMap: make(map[*registry.ServiceInstance]context.CancelFunc),
+		ctxMap: make(map[string]*serviceCancel),
 	}
 }
 
@@ -94,7 +99,10 @@ func (r *Registry) Register(ctx context.Context, service *registry.ServiceInstan
 	}
 
 	hctx, cancel := context.WithCancel(r.opts.ctx)
-	r.ctxMap[service] = cancel
+	r.ctxMap[service.ID] = &serviceCancel{
+		service: service,
+		cancel:  cancel,
+	}
 	go r.heartBeat(hctx, leaseID, key, value)
 	return nil
 }
@@ -107,9 +115,9 @@ func (r *Registry) Deregister(ctx context.Context, service *registry.ServiceInst
 		}
 	}()
 	// cancel heartbeat
-	if cancel, ok := r.ctxMap[service]; ok {
-		cancel()
-		delete(r.ctxMap, service)
+	if serviceCancel, ok := r.ctxMap[service.ID]; ok {
+		serviceCancel.cancel()
+		delete(r.ctxMap, service.ID)
 	}
 	key := fmt.Sprintf("%s/%s/%s", r.opts.namespace, service.Name, service.ID)
 	_, err := r.client.Delete(ctx, key)
diff --git a/contrib/registry/nacos/registry.go b/contrib/registry/nacos/registry.go
index 1a3722b8..40337691 100644
--- a/contrib/registry/nacos/registry.go
+++ b/contrib/registry/nacos/registry.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"net"
 	"net/url"
 	"strconv"
@@ -100,6 +101,7 @@ func (r *Registry) Register(_ context.Context, si *registry.ServiceInstance) err
 		if err != nil {
 			return err
 		}
+		weight := r.opts.weight
 		var rmd map[string]string
 		if si.Metadata == nil {
 			rmd = map[string]string{
@@ -113,12 +115,18 @@ func (r *Registry) Register(_ context.Context, si *registry.ServiceInstance) err
 			}
 			rmd["kind"] = u.Scheme
 			rmd["version"] = si.Version
+			if w, ok := si.Metadata["weight"]; ok {
+				weight, err = strconv.ParseFloat(w, 64)
+				if err != nil {
+					weight = r.opts.weight
+				}
+			}
 		}
 		_, e := r.cli.RegisterInstance(vo.RegisterInstanceParam{
 			Ip:          host,
 			Port:        uint64(p),
 			ServiceName: si.Name + "." + u.Scheme,
-			Weight:      r.opts.weight,
+			Weight:      weight,
 			Enable:      true,
 			Healthy:     true,
 			Ephemeral:   true,
@@ -180,16 +188,23 @@ func (r *Registry) GetService(_ context.Context, serviceName string) ([]*registr
 	items := make([]*registry.ServiceInstance, 0, len(res))
 	for _, in := range res {
 		kind := r.opts.kind
+		weight := r.opts.weight
 		if k, ok := in.Metadata["kind"]; ok {
 			kind = k
 		}
-		items = append(items, &registry.ServiceInstance{
+		if in.Weight > 0 {
+			weight = in.Weight
+		}
+
+		r := &registry.ServiceInstance{
 			ID:        in.InstanceId,
 			Name:      in.ServiceName,
 			Version:   in.Metadata["version"],
 			Metadata:  in.Metadata,
 			Endpoints: []string{fmt.Sprintf("%s://%s:%d", kind, in.Ip, in.Port)},
-		})
+		}
+		r.Metadata["weight"] = strconv.Itoa(int(math.Ceil(weight)))
+		items = append(items, r)
 	}
 	return items, nil
 }


Test patch:
diff --git a/contrib/registry/nacos/registry_test.go b/contrib/registry/nacos/registry_test.go
index 95768eea..bfa7b001 100644
--- a/contrib/registry/nacos/registry_test.go
+++ b/contrib/registry/nacos/registry_test.go
@@ -13,10 +13,12 @@ import (
 	"github.com/go-kratos/kratos/v2/registry"
 )
 
+var testServerConfig = []constant.ServerConfig{
+	*constant.NewServerConfig("127.0.0.1", 8848),
+}
+
 func TestRegistry_Register(t *testing.T) {
-	sc := []constant.ServerConfig{
-		*constant.NewServerConfig("127.0.0.1", 8848),
-	}
+	sc := testServerConfig
 
 	cc := constant.ClientConfig{
 		NamespaceId:         "public", // namespace id
@@ -231,9 +233,7 @@ func TestRegistry_Deregister(t *testing.T) {
 			},
 			wantErr: false,
 			preFunc: func(t *testing.T) {
-				sc := []constant.ServerConfig{
-					*constant.NewServerConfig("127.0.0.1", 8848),
-				}
+				sc := testServerConfig
 
 				cc := constant.ClientConfig{
 					NamespaceId:         "public", // namespace id
@@ -292,9 +292,7 @@ func TestRegistry_Deregister(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			sc := []constant.ServerConfig{
-				*constant.NewServerConfig("127.0.0.1", 8848),
-			}
+			sc := testServerConfig
 
 			cc := constant.ClientConfig{
 				NamespaceId:         "public", // namespace id
@@ -329,9 +327,7 @@ func TestRegistry_Deregister(t *testing.T) {
 }
 
 func TestRegistry_GetService(t *testing.T) {
-	sc := []constant.ServerConfig{
-		*constant.NewServerConfig("127.0.0.1", 8848),
-	}
+	sc := testServerConfig
 
 	cc := constant.ClientConfig{
 		NamespaceId:         "public", // namespace id
@@ -385,7 +381,7 @@ func TestRegistry_GetService(t *testing.T) {
 				if err != nil {
 					t.Error(err)
 				}
-				time.Sleep(time.Second)
+				time.Sleep(time.Second * 3)
 			},
 			deferFunc: func(t *testing.T) {
 				err = r.Deregister(context.Background(), testServer)
@@ -404,7 +400,7 @@ func TestRegistry_GetService(t *testing.T) {
 				ID:        "127.0.0.1#8080#DEFAULT#DEFAULT_GROUP@@test3.grpc",
 				Name:      "DEFAULT_GROUP@@test3.grpc",
 				Version:   "v1.0.0",
-				Metadata:  map[string]string{"version": "v1.0.0", "kind": "grpc"},
+				Metadata:  map[string]string{"version": "v1.0.0", "kind": "grpc", "weight": "100"},
 				Endpoints: []string{"grpc://127.0.0.1:8080"},
 			}},
 			wantErr: false,
@@ -445,9 +441,7 @@ func TestRegistry_GetService(t *testing.T) {
 }
 
 func TestRegistry_Watch(t *testing.T) {
-	sc := []constant.ServerConfig{
-		*constant.NewServerConfig("127.0.0.1", 8848),
-	}
+	sc := testServerConfig
 
 	cc := constant.ClientConfig{
 		NamespaceId:         "public", // namespace id