support grpc health check (#2493)

* support grpc health check protocol * add grpc health check ut --------- Co-authored-by: jiangyuting <jiangyutingwangyi@163.com>
apache · Jan 23, 2024 · ff75449 · ff75449
1 parent 7dbe775
commit ff75449
Show file tree

Hide file tree

Showing 6 changed files with 194 additions and 0 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -456,6 +456,7 @@ set(PROTO_FILES idl_options.proto
                 brpc/get_favicon.proto
                 brpc/span.proto
                 brpc/builtin_service.proto
+                brpc/grpc_health_check.proto
                 brpc/get_js.proto
                 brpc/errno.proto
                 brpc/nshead_meta.proto

diff --git a/src/brpc/builtin/grpc_health_check_service.cpp b/src/brpc/builtin/grpc_health_check_service.cpp
@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "brpc/builtin/common.h"
+#include "brpc/builtin/grpc_health_check_service.h"
+#include "brpc/controller.h"           // Controller
+#include "brpc/closure_guard.h"        // ClosureGuard
+#include "brpc/server.h"               // Server
+
+namespace brpc {
+    void GrpcHealthCheckService::Check(::google::protobuf::RpcController* cntl_base,
+                                  const grpc::health::v1::HealthCheckRequest* request,
+                                  grpc::health::v1::HealthCheckResponse* response,
+                                  ::google::protobuf::Closure* done) {
+        ClosureGuard done_guard(done);
+        Controller *cntl = static_cast<Controller*>(cntl_base);
+        const Server* server = cntl->server();
+        if (server->options().health_reporter) {
+            server->options().health_reporter->GenerateReport(
+                    cntl, done_guard.release());
+        } else {
+            response->set_status(grpc::health::v1::HealthCheckResponse_ServingStatus_SERVING);
+        }
+    }
+
+} // namespace brpc
+
diff --git a/src/brpc/builtin/grpc_health_check_service.h b/src/brpc/builtin/grpc_health_check_service.h
@@ -0,0 +1,37 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+
+#ifndef BRPC_GRPC_HEALTH_CHECK_SERVICE_H
+#define BRPC_GRPC_HEALTH_CHECK_SERVICE_H
+
+#include "brpc/grpc_health_check.pb.h"
+
+namespace brpc {
+
+class GrpcHealthCheckService : public grpc::health::v1::Health {
+public:
+    void Check(::google::protobuf::RpcController* cntl_base,
+               const grpc::health::v1::HealthCheckRequest* request,
+               grpc::health::v1::HealthCheckResponse* response,
+               ::google::protobuf::Closure* done);
+
+};
+
+} // namespace brpc
+
+#endif //BRPC_GRPC_HEALTH_CHECK_SERVICE_H
diff --git a/src/brpc/grpc_health_check.proto b/src/brpc/grpc_health_check.proto
@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+syntax = "proto2";
+package grpc.health.v1;
+option cc_generic_services = true;
+
+message HealthCheckRequest {
+    optional string service = 1;
+}
+
+message HealthCheckResponse {
+    enum ServingStatus {
+        UNKNOWN = 0;
+        SERVING = 1;
+        NOT_SERVING = 2;
+        SERVICE_UNKNOWN = 3;  // Used only by the Watch method.
+    }
+    optional ServingStatus status = 1;
+}
+
+service Health {
+    rpc Check(HealthCheckRequest) returns (HealthCheckResponse);
+
+    rpc Watch(HealthCheckRequest) returns (stream HealthCheckResponse);
+}
+
diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp
@@ -45,6 +45,7 @@
 #include "brpc/builtin/bad_method_service.h"   // BadMethodService
 #include "brpc/builtin/get_favicon_service.h"
 #include "brpc/builtin/get_js_service.h"
+#include "brpc/builtin/grpc_health_check_service.h"  // GrpcHealthCheckService
 #include "brpc/builtin/version_service.h"
 #include "brpc/builtin/health_service.h"
 #include "brpc/builtin/list_service.h"
@@ -564,6 +565,10 @@ int Server::AddBuiltinServices() {
         LOG(ERROR) << "Fail to add GetJsService";
         return -1;
     }
+    if (AddBuiltinService(new (std::nothrow) GrpcHealthCheckService)) {
+        LOG(ERROR) << "Fail to add GrpcHealthCheckService";
+        return -1;
+    }
     return 0;
 }
 

diff --git a/test/brpc_builtin_service_unittest.cpp b/test/brpc_builtin_service_unittest.cpp
@@ -53,6 +53,8 @@
 #include "brpc/builtin/common.h"
 #include "brpc/builtin/bad_method_service.h"
 #include "echo.pb.h"
+#include "brpc/grpc_health_check.pb.h"
+#include "json2pb/pb_to_json.h"
 
 DEFINE_bool(foo, false, "Flags for UT");
 BRPC_VALIDATE_GFLAG(foo, brpc::PassValidate);
@@ -552,6 +554,73 @@ TEST_F(BuiltinServiceTest, customized_health) {
     EXPECT_EQ("i'm ok", cntl.response_attachment());
 }
 
+class MyGrpcHealthReporter : public brpc::HealthReporter {
+public:
+    void GenerateReport(brpc::Controller* cntl,
+                        google::protobuf::Closure* done) {
+        grpc::health::v1::HealthCheckResponse response;
+        response.set_status(grpc::health::v1::HealthCheckResponse_ServingStatus_UNKNOWN);
+
+        if (cntl->response()) {
+            cntl->response()->CopyFrom(response);
+        } else {
+            std::string json;
+            json2pb::ProtoMessageToJson(response, &json);
+            cntl->http_response().set_content_type("application/json");
+            cntl->response_attachment().append(json);
+        }
+        done->Run();
+    }
+};
+
+TEST_F(BuiltinServiceTest, normal_grpc_health) {
+    brpc::ServerOptions opt;
+    ASSERT_EQ(0, _server.Start(9798, &opt));
+
+    grpc::health::v1::HealthCheckResponse response;
+    grpc::health::v1::HealthCheckRequest request;
+    request.set_service("grpc_req_from_brpc");
+    brpc::Controller cntl;
+    brpc::ChannelOptions copt;
+    copt.protocol = "h2:grpc";
+    brpc::Channel chan;
+    ASSERT_EQ(0, chan.Init("127.0.0.1:9798", &copt));
+    grpc::health::v1::Health_Stub stub(&chan);
+    stub.Check(&cntl, &request, &response, NULL);
+    EXPECT_FALSE(cntl.Failed()) << cntl.ErrorText();
+    EXPECT_EQ(response.status(), grpc::health::v1::HealthCheckResponse_ServingStatus_SERVING);
+
+    response.Clear();
+    brpc::Controller cntl1;
+    cntl1.http_request().uri() = "/grpc.health.v1.Health/Check";
+    chan.CallMethod(NULL, &cntl1, &request, &response, NULL);
+    EXPECT_FALSE(cntl.Failed()) << cntl.ErrorText();
+    EXPECT_EQ(response.status(), grpc::health::v1::HealthCheckResponse_ServingStatus_SERVING);
+}
+
+TEST_F(BuiltinServiceTest, customized_grpc_health) {
+    brpc::ServerOptions opt;
+    MyGrpcHealthReporter hr;
+    opt.health_reporter = &hr;
+    ASSERT_EQ(0, _server.Start(9798, &opt));
+
+    grpc::health::v1::HealthCheckResponse response;
+    grpc::health::v1::HealthCheckRequest request;
+    request.set_service("grpc_req_from_brpc");
+    brpc::Controller cntl;
+
+    brpc::ChannelOptions copt;
+    copt.protocol = "h2:grpc";
+    brpc::Channel chan;
+    ASSERT_EQ(0, chan.Init("127.0.0.1:9798", &copt));
+
+    grpc::health::v1::Health_Stub stub(&chan);
+    stub.Check(&cntl, &request, &response, NULL);
+
+    EXPECT_FALSE(cntl.Failed()) << cntl.ErrorText();
+    EXPECT_EQ(response.status(), grpc::health::v1::HealthCheckResponse_ServingStatus_UNKNOWN);
+}
+
 TEST_F(BuiltinServiceTest, status) {
     TestStatus(false);
     TestStatus(true);