From 88874e729858c71deab1598257f903e482618c0f Mon Sep 17 00:00:00 2001
From: EvanYang <yangyi2048@outlook.com>
Date: Tue, 19 Dec 2023 16:25:18 +0800
Subject: [PATCH] feat: add custom-healthcheck

---
 custom-healthcheck/main.go | 69 ++++++++++++++++++++++++++++++++++++++
 custom-healthcheck/test.go | 43 ++++++++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 custom-healthcheck/main.go
 create mode 100644 custom-healthcheck/test.go

diff --git a/custom-healthcheck/main.go b/custom-healthcheck/main.go
new file mode 100644
index 0000000..746edcb
--- /dev/null
+++ b/custom-healthcheck/main.go
@@ -0,0 +1,69 @@
+package main
+
+import (
+	"context"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/status"
+	"os"
+
+	"github.com/go-kratos/kratos/v2"
+	"github.com/go-kratos/kratos/v2/log"
+	"github.com/go-kratos/kratos/v2/transport/grpc"
+)
+
+type HealthCheckService struct {
+	log *log.Helper
+}
+
+func NewHealthCheckService(logger log.Logger) *HealthCheckService {
+	return &HealthCheckService{
+		log: log.NewHelper(logger),
+	}
+}
+
+func (s *HealthCheckService) Check(ctx context.Context, req *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
+	s.log.WithContext(ctx).Infof("HealthCheckService.Check: %v", req)
+
+	// TODO: add your health check logic here
+
+	return &grpc_health_v1.HealthCheckResponse{
+		Status: grpc_health_v1.HealthCheckResponse_SERVING,
+	}, nil
+}
+
+func (s *HealthCheckService) Watch(req *grpc_health_v1.HealthCheckRequest, srv grpc_health_v1.Health_WatchServer) error {
+	s.log.Infof("HealthCheckService.Watch: %v, %v", req, srv)
+
+	// simple implementation
+	resp := &grpc_health_v1.HealthCheckResponse{
+		Status: grpc_health_v1.HealthCheckResponse_SERVING,
+	}
+	if err := srv.Send(resp); err != nil {
+		return status.Errorf(codes.Unknown, "failed to send response: %v", err)
+	}
+	return nil
+}
+
+func main() {
+	grpcSrv := grpc.NewServer(
+		grpc.Address(":9000"),
+		grpc.Middleware(),
+		grpc.CustomHealth(),
+	)
+	app := kratos.New(
+		kratos.Name("helloworld"),
+		kratos.Server(
+			grpcSrv,
+		),
+	)
+	logger := log.With(log.NewStdLogger(os.Stdout),
+		"ts", log.DefaultTimestamp,
+		"caller", log.DefaultCaller,
+	)
+	health := NewHealthCheckService(logger)
+	grpc_health_v1.RegisterHealthServer(grpcSrv, health)
+	if err := app.Run(); err != nil {
+		log.Fatal(err)
+	}
+}
diff --git a/custom-healthcheck/test.go b/custom-healthcheck/test.go
new file mode 100644
index 0000000..92db78c
--- /dev/null
+++ b/custom-healthcheck/test.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"context"
+	"github.com/go-kratos/kratos/v2/log"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health/grpc_health_v1"
+)
+
+func main() {
+	// init grpc conn
+	conn, err := grpc.Dial(":9000", grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		panic(err)
+	}
+	healthClient := grpc_health_v1.NewHealthClient(conn)
+	checkReq := &grpc_health_v1.HealthCheckRequest{
+		Service: "helloworld.Greeter",
+	}
+	ctx := context.Background()
+	checkResp, err := healthClient.Check(ctx, checkReq)
+	if err != nil {
+		panic(err)
+	}
+	log.Info("checkResp: ", checkResp)
+
+	watchReq := &grpc_health_v1.HealthCheckRequest{
+		Service: "helloworld.Greeter",
+	}
+	watchResp, err := healthClient.Watch(ctx, watchReq)
+	if err != nil {
+		panic(err)
+	}
+	// loop watch
+	for {
+		resp, err := watchResp.Recv()
+		if err != nil {
+			panic(err)
+		}
+		log.Info("watchResp: ", resp)
+	}
+}