Repository: go-kratos/kratos (go)

Issue:
[Question] Why does grpc server not open the configuration of health Serveroption?
我们期望在 GRPC Server 的健康检查中实现一些自定义的逻辑，用以更加可靠地实现 health， ready 探测，但发现没有提供相应的 ServerOption，以致无法注入自定义的逻辑。

请问：出于什么考虑，没有通过 ServerOption 开放对 health 的配置？

https://github.com/go-kratos/kratos/blob/main/transport/grpc/server.go#L118

Pull request:
feat(transport/grpc): support custom healthcheck
#### Description (what this PR does / why we need it):
Currently, it is not possible to define custom health checks. This PR should solve that.


#### Which issue(s) this PR fixes (resolves / be part of):
resolves #1397 


Patch:
diff --git a/contrib/config/apollo/apollo.go b/contrib/config/apollo/apollo.go
index bd96ef11..d463df29 100644
--- a/contrib/config/apollo/apollo.go
+++ b/contrib/config/apollo/apollo.go
@@ -25,6 +25,8 @@ const (
 	properties = "properties"
 )
 
+var formats map[string]struct{}
+
 // Option is apollo option
 type Option func(*options)
 
@@ -129,10 +131,16 @@ func NewSource(opts ...Option) config.Source {
 
 func format(ns string) string {
 	arr := strings.Split(ns, ".")
-	if len(arr) <= 1 || arr[len(arr)-1] == properties {
+	suffix := arr[len(arr)-1]
+	if len(arr) <= 1 || suffix == properties {
+		return json
+	}
+	if _, ok := formats[suffix]; !ok {
+		// fallback
 		return json
 	}
-	return arr[len(arr)-1]
+
+	return suffix
 }
 
 func (e *apollo) load() []*config.KeyValue {
@@ -149,7 +157,7 @@ func (e *apollo) load() []*config.KeyValue {
 			kvs = append(kvs, kv)
 			continue
 		}
-		if strings.Contains(ns, ".") && !strings.Contains(ns, properties) &&
+		if strings.Contains(ns, ".") && !strings.HasSuffix(ns, "."+properties) &&
 			(format(ns) == yaml || format(ns) == yml || format(ns) == json) {
 			kv, err := e.getOriginConfig(ns)
 			if err != nil {
@@ -252,16 +260,28 @@ func resolve(key string, value interface{}, target map[string]interface{}) {
 // eg: namespace.ext with subKey got namespace.subKey
 func genKey(ns, sub string) string {
 	arr := strings.Split(ns, ".")
-	if len(arr) < 1 {
-		return sub
-	}
-
 	if len(arr) == 1 {
 		if ns == "" {
 			return sub
 		}
+
 		return ns + "." + sub
 	}
 
-	return strings.Join(arr[:len(arr)-1], ".") + "." + sub
+	suffix := arr[len(arr)-1]
+	_, ok := formats[suffix]
+	if ok {
+		return strings.Join(arr[:len(arr)-1], ".") + "." + sub
+	}
+
+	return ns + "." + sub
+}
+
+func init() {
+	formats = make(map[string]struct{})
+
+	formats[yaml] = struct{}{}
+	formats[yml] = struct{}{}
+	formats[json] = struct{}{}
+	formats[properties] = struct{}{}
 }
diff --git a/contrib/config/apollo/watcher.go b/contrib/config/apollo/watcher.go
index ea5ef108..d59bdd21 100644
--- a/contrib/config/apollo/watcher.go
+++ b/contrib/config/apollo/watcher.go
@@ -25,7 +25,7 @@ type customChangeListener struct {
 
 func (c *customChangeListener) onChange(namespace string, changes map[string]*storage.ConfigChange) []*config.KeyValue {
 	kv := make([]*config.KeyValue, 0, 2)
-	if strings.Contains(namespace, ".") && !strings.Contains(namespace, properties) &&
+	if strings.Contains(namespace, ".") && !strings.HasSuffix(namespace, "."+properties) &&
 		(format(namespace) == yaml || format(namespace) == yml || format(namespace) == json) {
 		value, err := c.apollo.client.GetConfigCache(namespace).Get("content")
 		if err != nil {
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index 9e5ff17f..e3a0cea9 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -67,6 +67,13 @@ func Middleware(m ...middleware.Middleware) ServerOption {
 	}
 }
 
+// CustomHealth Checks server.
+func CustomHealth() ServerOption {
+	return func(s *Server) {
+		s.customHealth = true
+	}
+}
+
 // TLSConfig with TLS config.
 func TLSConfig(c *tls.Config) ServerOption {
 	return func(s *Server) {
@@ -105,20 +112,21 @@ func Options(opts ...grpc.ServerOption) ServerOption {
 // Server is a gRPC server wrapper.
 type Server struct {
 	*grpc.Server
-	baseCtx    context.Context
-	tlsConf    *tls.Config
-	lis        net.Listener
-	err        error
-	network    string
-	address    string
-	endpoint   *url.URL
-	timeout    time.Duration
-	middleware matcher.Matcher
-	unaryInts  []grpc.UnaryServerInterceptor
-	streamInts []grpc.StreamServerInterceptor
-	grpcOpts   []grpc.ServerOption
-	health     *health.Server
-	metadata   *apimd.Server
+	baseCtx      context.Context
+	tlsConf      *tls.Config
+	lis          net.Listener
+	err          error
+	network      string
+	address      string
+	endpoint     *url.URL
+	timeout      time.Duration
+	middleware   matcher.Matcher
+	unaryInts    []grpc.UnaryServerInterceptor
+	streamInts   []grpc.StreamServerInterceptor
+	grpcOpts     []grpc.ServerOption
+	health       *health.Server
+	customHealth bool
+	metadata     *apimd.Server
 }
 
 // NewServer creates a gRPC server by options.
@@ -159,7 +167,9 @@ func NewServer(opts ...ServerOption) *Server {
 	srv.Server = grpc.NewServer(grpcOpts...)
 	srv.metadata = apimd.NewServer(srv.Server)
 	// internal register
-	grpc_health_v1.RegisterHealthServer(srv.Server, srv.health)
+	if !srv.customHealth {
+		grpc_health_v1.RegisterHealthServer(srv.Server, srv.health)
+	}
 	apimd.RegisterMetadataServer(srv.Server, srv.metadata)
 	reflection.Register(srv.Server)
 	return srv


Test patch:
diff --git a/contrib/config/apollo/apollo_test.go b/contrib/config/apollo/apollo_test.go
new file mode 100644
index 00000000..185ca4c6
--- /dev/null
+++ b/contrib/config/apollo/apollo_test.go
@@ -0,0 +1,90 @@
+package apollo
+
+import (
+	"testing"
+)
+
+func Test_genKey(t *testing.T) {
+	type args struct {
+		ns  string
+		sub string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "blank namespace",
+			args: args{
+				ns:  "",
+				sub: "x.y",
+			},
+			want: "x.y",
+		},
+		{
+			name: "properties namespace",
+			args: args{
+				ns:  "application",
+				sub: "x.y",
+			},
+			want: "application.x.y",
+		},
+		{
+			name: "namespace with format",
+			args: args{
+				ns:  "app.yaml",
+				sub: "x.y",
+			},
+			want: "app.x.y",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := genKey(tt.args.ns, tt.args.sub); got != tt.want {
+				t.Errorf("genKey() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func Test_format(t *testing.T) {
+	tests := []struct {
+		name      string
+		namespace string
+		want      string
+	}{
+		{
+			name:      "properties namespace",
+			namespace: "application",
+			want:      "json",
+		},
+		{
+			name:      "properties namespace #1",
+			namespace: "app.setting",
+			want:      "json",
+		},
+		{
+			name:      "namespace with format[yaml]",
+			namespace: "app.yaml",
+			want:      "yaml",
+		},
+		{
+			name:      "namespace with format[yml]",
+			namespace: "app.yml",
+			want:      "yml",
+		},
+		{
+			name:      "namespace with format[json]",
+			namespace: "app.json",
+			want:      "json",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := format(tt.namespace); got != tt.want {
+				t.Errorf("format() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}