Merge pull request #573 from zyearn/discovery

add discovery client

src/brpc/policy/discovery_naming_service.h

@@ -19,10 +19,65 @@
 
 #include "brpc/periodic_naming_service.h"
 #include "brpc/channel.h"
+#include "butil/synchronization/lock.h"
 
 namespace brpc {
 namespace policy {
 
+struct DiscoveryRegisterParam {
+    std::string appid;
+    std::string hostname;
+    std::string env;
+    std::string zone;
+    std::string region;
+    std::string addrs;          // splitted by ','
+    int status;
+    std::string version;
+    std::string metadata;
+
+    bool IsValid() const;
+};
+
+struct DiscoveryFetchsParam {
+    std::string appid;
+    std::string env;
+    std::string status;
+
+    bool IsValid() const;
+};
+
+// ONE DiscoveryClient corresponds to ONE service instance.
+// If your program has multiple service instances to register,
+// you need multiple DiscoveryClient.
+// Note: Unregister is automatically called in dtor.
+class DiscoveryClient {
+public:
+    DiscoveryClient();
+    ~DiscoveryClient();
+
+    int Register(const DiscoveryRegisterParam& req);
+    int Fetchs(const DiscoveryFetchsParam& req, std::vector<ServerNode>* servers) const;
+
+private:
+    static void* PeriodicRenew(void* arg);
+    int DoCancel() const;
+    int DoRegister() const;
+    int DoRenew() const;
+
+private:
+    bthread_t _th;
+    butil::atomic<bool> _registered;
+    std::string _appid;
+    std::string _hostname;
+    std::string _addrs;
+    std::string _env;
+    std::string _region;
+    std::string _zone;
+    int _status;
+    std::string _version;
+    std::string _metadata;
+};
+
 class DiscoveryNamingService : public PeriodicNamingService {
 private:
     int GetServers(const char* service_name,
@@ -35,16 +90,11 @@ private:
     void Destroy() override;
 
 private:
-    int ParseNodesResult(const butil::IOBuf& buf, std::string* server_addr);
-    int ParseFetchsResult(const butil::IOBuf& buf, const char* service_name,
-                            std::vector<ServerNode>* servers);
-
-    Channel _channel;
-    bool _is_initialized = false;
+    DiscoveryClient _client;
 };
 
+
 } // namespace policy
 } // namespace brpc
 
-
 #endif // BRPC_POLICY_DISCOVERY_NAMING_SERVICE_H

src/brpc/socket.cpp

@@ -863,7 +863,7 @@ int Socket::SetFailed(int error_code, const char* error_fmt, ...) {
             // by Channel to revive never-connected socket when server side
             // comes online.
             if (_health_check_interval_s > 0) {
-                GetOrNewSharedPart( )->circuit_breaker.MarkAsBroken();
+                GetOrNewSharedPart()->circuit_breaker.MarkAsBroken();
                 PeriodicTaskManager::StartTaskAt(
                     new HealthCheckTask(id()),
                     butil::milliseconds_from_now(GetOrNewSharedPart()->

test/brpc_naming_service_unittest.cpp

@@ -30,6 +30,13 @@ DECLARE_string(consul_file_naming_service_dir);
 DECLARE_string(consul_service_discovery_url);
 DECLARE_string(discovery_api_addr);
 DECLARE_string(discovery_env);
+DECLARE_int32(discovery_renew_interval_s);
+
+// Defined in discovery_naming_service.cpp
+int ParseFetchsResult(const butil::IOBuf& buf,
+                      const char* service_name,
+                      std::vector<brpc::ServerNode>* servers);
+int ParseNodesResult(const butil::IOBuf& buf, std::string* server_addr);
 
 } // policy
 } // brpc
@@ -445,7 +452,7 @@ static const std::string s_fetchs_result = R"({
                     },
                     "addrs":[
                         "http://127.0.0.1:8999",
-                        "gorpc://127.0.1.1:9000"
+                        "grpc://127.0.1.1:9000"
                     ],
                     "status":1,
                     "reg_timestamp":1539001034551496412,
@@ -472,7 +479,7 @@ static const std::string s_fetchs_result = R"({
                         },
                         "addrs":[
                             "http://127.0.0.1:8999",
-                            "gorpc://127.0.1.1:9000"
+                            "grpc://127.0.1.1:9000"
                         ],
                         "status":1,
                         "reg_timestamp":1539001034551496412,
@@ -510,23 +517,26 @@ static std::string s_nodes_result = R"({
     ]
 })";
 
+
 TEST(NamingServiceTest, discovery_parse_function) {
     std::vector<brpc::ServerNode> servers;
     brpc::policy::DiscoveryNamingService dcns;
     butil::IOBuf buf;
     buf.append(s_fetchs_result);
-    ASSERT_EQ(0, dcns.ParseFetchsResult(buf, "admin.test", &servers));
-    ASSERT_EQ((size_t)2, servers.size());
+    ASSERT_EQ(0, brpc::policy::ParseFetchsResult(buf, "admin.test", &servers));
+    ASSERT_EQ((size_t)1, servers.size());
     buf.clear();
     buf.append(s_nodes_result);
     std::string server;
-    ASSERT_EQ(0, dcns.ParseNodesResult(buf, &server));
+    ASSERT_EQ(0, brpc::policy::ParseNodesResult(buf, &server));
     ASSERT_EQ("127.0.0.1:8635", server);
 }
 
 class DiscoveryNamingServiceImpl : public test::DiscoveryNamingService {
 public:
-    DiscoveryNamingServiceImpl () {}
+    DiscoveryNamingServiceImpl()
+        : _renew_count(0)
+        , _cancel_count(0) {}
     virtual ~DiscoveryNamingServiceImpl() {}
 
     void Nodes(google::protobuf::RpcController* cntl_base,
@@ -546,15 +556,67 @@ public:
         brpc::Controller* cntl = static_cast<brpc::Controller*>(cntl_base);
         cntl->response_attachment().append(s_fetchs_result);
     }
+
+    void Register(google::protobuf::RpcController* cntl_base,
+                 const test::HttpRequest*,
+                 test::HttpResponse*,
+                 google::protobuf::Closure* done) {
+        brpc::ClosureGuard done_guard(done);
+        brpc::Controller* cntl = static_cast<brpc::Controller*>(cntl_base);
+        cntl->response_attachment().append(R"({
+            "code": 0,
+            "message": "0"
+        })");
+        return;
+    }
+
+    void Renew(google::protobuf::RpcController* cntl_base,
+               const test::HttpRequest*,
+               test::HttpResponse*,
+               google::protobuf::Closure* done) {
+        brpc::ClosureGuard done_guard(done);
+        brpc::Controller* cntl = static_cast<brpc::Controller*>(cntl_base);
+        cntl->response_attachment().append(R"({
+            "code": 0,
+            "message": "0"
+        })");
+        _renew_count++;
+        return;
+    }
+
+    void Cancel(google::protobuf::RpcController* cntl_base,
+                const test::HttpRequest*,
+                test::HttpResponse*,
+                google::protobuf::Closure* done) {
+        brpc::ClosureGuard done_guard(done);
+        brpc::Controller* cntl = static_cast<brpc::Controller*>(cntl_base);
+        cntl->response_attachment().append(R"({
+            "code": 0,
+            "message": "0"
+        })");
+        _cancel_count++;
+        return;
+    }
+
+    int RenewCount() const { return _renew_count; }
+    int CancelCount() const { return _cancel_count; }
+
+private:
+    int _renew_count;
+    int _cancel_count;
 };
 
 TEST(NamingServiceTest, discovery_sanity) {
     brpc::policy::FLAGS_discovery_api_addr = "http://127.0.0.1:8635/discovery/nodes";
+    brpc::policy::FLAGS_discovery_renew_interval_s = 1;
     brpc::Server server;
     DiscoveryNamingServiceImpl svc;
     std::string rest_mapping =
         "/discovery/nodes => Nodes, "
-        "/discovery/fetchs => Fetchs";
+        "/discovery/fetchs => Fetchs, "
+        "/discovery/register => Register, "
+        "/discovery/renew => Renew, "
+        "/discovery/cancel => Cancel";
     ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE,
                 rest_mapping.c_str()));
     ASSERT_EQ(0, server.Start("localhost:8635", NULL));
@@ -562,8 +624,31 @@ TEST(NamingServiceTest, discovery_sanity) {
     brpc::policy::DiscoveryNamingService dcns;
     std::vector<brpc::ServerNode> servers;
     ASSERT_EQ(0, dcns.GetServers("admin.test", &servers));
-    ASSERT_EQ((size_t)2, servers.size());
+    ASSERT_EQ((size_t)1, servers.size());
+
+    brpc::policy::DiscoveryRegisterParam dparam;
+    dparam.appid = "main.test";
+    dparam.hostname = "hostname";
+    dparam.addrs = "grpc://10.0.0.1:8000";
+    dparam.env = "dev";
+    dparam.zone = "sh001";
+    dparam.status = 1;
+    dparam.version = "v1";
+    {
+        brpc::policy::DiscoveryClient dc;
+    }
+    // Cancel is called iff Register is called
+    ASSERT_EQ(svc.CancelCount(), 0);
+    {
+        brpc::policy::DiscoveryClient dc;
+        // Two Register should start one Renew task , and make
+        // svc.RenewCount() be one.
+        ASSERT_EQ(0, dc.Register(dparam));
+        ASSERT_EQ(0, dc.Register(dparam));
+        bthread_usleep(1000000);
+    }
+    ASSERT_EQ(svc.RenewCount(), 1);
+    ASSERT_EQ(svc.CancelCount(), 1);
 }
 
-
 } //namespace

test/bthread_butex_unittest.cpp

@@ -208,7 +208,7 @@ TEST(ButexTest, wait_without_stop) {
         ASSERT_EQ(0, bthread_join(th, NULL));
         tm.stop();
         
-        ASSERT_LT(labs(tm.m_elapsed() - WAIT_MSEC), 40);
+        ASSERT_LT(labs(tm.m_elapsed() - WAIT_MSEC), 250);
     }
     bthread::butex_destroy(butex);
 }

test/echo.proto

@@ -56,6 +56,9 @@ service UserNamingService {
 service DiscoveryNamingService {
     rpc Nodes(HttpRequest) returns (HttpResponse);
     rpc Fetchs(HttpRequest) returns (HttpResponse);
+    rpc Register(HttpRequest) returns (HttpResponse);
+    rpc Renew(HttpRequest) returns (HttpResponse);
+    rpc Cancel(HttpRequest) returns (HttpResponse);
 };
 
 enum State0 {