fix whitelist disabled case, update ut

vesoft-inc · May 30, 2019 · 2fbb60a · 2fbb60a
1 parent 3c72ef0
commit 2fbb60a
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 7 deletions.
diff --git a/src/meta/processors/BaseProcessor.inl b/src/meta/processors/BaseProcessor.inl
@@ -8,6 +8,8 @@
 #include "meta/processors/BaseProcessor.h"
 #include "meta/ActiveHostsMan.h"
 
+DECLARE_bool(hosts_whitelist_enabled);
+
 namespace nebula {
 namespace meta {
 
@@ -130,13 +132,6 @@ StatusOr<std::vector<nebula::cpp2::HostAddr>> BaseProcessor<RESP>::allHosts() {
 template<typename RESP>
 StatusOr<std::vector<cpp2::HostItem>> BaseProcessor<RESP>::allHostsWithStatus() {
     std::vector<cpp2::HostItem> hostItems;
-    const auto& prefix = MetaServiceUtils::hostPrefix();
-    std::unique_ptr<kvstore::KVIterator> iter;
-    auto ret = kvstore_->prefix(kDefaultSpaceId_, kDefaultPartId_, prefix, &iter);
-    if (ret != kvstore::ResultCode::SUCCEEDED) {
-        return Status::Error("Can't find any hosts");
-    }
-
     auto hosts = ActiveHostsManHolder::hostsMan()->getActiveHosts();
     std::vector<nebula::cpp2::HostAddr> activeHosts;
     activeHosts.resize(hosts.size());
@@ -147,6 +142,23 @@ StatusOr<std::vector<cpp2::HostItem>> BaseProcessor<RESP>::allHostsWithStatus()
         return th;
     });
 
+    if (!FLAGS_hosts_whitelist_enabled) {
+        std::for_each(activeHosts.begin(), activeHosts.end(), [&](auto& h) {
+            cpp2::HostItem item;
+            item.set_hostAddr(h);
+            item.set_status("online");
+            hostItems.emplace_back(item);
+        });
+        return hostItems;
+    }
+
+    const auto& prefix = MetaServiceUtils::hostPrefix();
+    std::unique_ptr<kvstore::KVIterator> iter;
+    auto ret = kvstore_->prefix(kDefaultSpaceId_, kDefaultPartId_, prefix, &iter);
+    if (ret != kvstore::ResultCode::SUCCEEDED) {
+        return Status::Error("Can't find any hosts");
+    }
+
     while (iter->valid()) {
         cpp2::HostItem item;
         nebula::cpp2::HostAddr host;

diff --git a/src/meta/test/ProcessorTest.cpp b/src/meta/test/ProcessorTest.cpp
@@ -40,8 +40,11 @@ namespace meta {
 using nebula::cpp2::SupportedType;
 using apache::thrift::FragileConstructor::FRAGILE;
 
+
 TEST(ProcessorTest, AddHostsTest) {
     fs::TempDir rootPath("/tmp/AddHostsTest.XXXXXX");
+    FLAGS_expired_hosts_check_interval_sec = 1;
+    FLAGS_expired_threshold_sec = 1;
     std::unique_ptr<kvstore::KVStore> kv(TestUtils::initKV(rootPath.path()));
     {
         std::vector<nebula::cpp2::HostAddr> thriftHosts;
@@ -116,6 +119,55 @@ TEST(ProcessorTest, AddHostsTest) {
     }
 }
 
+TEST(ProcessorTest, ListHostsTest) {
+    fs::TempDir rootPath("/tmp/ListHostsTest.XXXXXX");
+    std::unique_ptr<kvstore::KVStore> kv(TestUtils::initKV(rootPath.path()));
+    {
+        std::vector<nebula::cpp2::HostAddr> thriftHosts;
+        std::vector<HostAddr> hosts;
+        for (auto i = 0; i < 10; i++) {
+            thriftHosts.emplace_back(FRAGILE, i, i);
+            hosts.emplace_back(i, i);
+        }
+        meta::TestUtils::registerHB(hosts);
+
+        cpp2::AddHostsReq req;
+        req.set_hosts(std::move(thriftHosts));
+        auto* processor = AddHostsProcessor::instance(kv.get());
+        auto f = processor->getFuture();
+        processor->process(req);
+        auto resp = std::move(f).get();
+        ASSERT_EQ(cpp2::ErrorCode::SUCCEEDED, resp.code);
+    }
+    {
+        cpp2::ListHostsReq req;
+        auto* processor = ListHostsProcessor::instance(kv.get());
+        auto f = processor->getFuture();
+        processor->process(req);
+        auto resp = std::move(f).get();
+        ASSERT_EQ(10, resp.hosts.size());
+        for (auto i = 0; i < 10; i++) {
+            ASSERT_EQ(i, resp.hosts[i].hostAddr.ip);
+            ASSERT_EQ(i, resp.hosts[i].hostAddr.port);
+            ASSERT_EQ("online", resp.hosts[i].status);
+        }
+    }
+    {
+        sleep(FLAGS_expired_threshold_sec + 1);
+        cpp2::ListHostsReq req;
+        auto* processor = ListHostsProcessor::instance(kv.get());
+        auto f = processor->getFuture();
+        processor->process(req);
+        auto resp = std::move(f).get();
+        ASSERT_EQ(10, resp.hosts.size());
+        for (auto i = 0; i < 10; i++) {
+            ASSERT_EQ(i, resp.hosts[i].hostAddr.ip);
+            ASSERT_EQ(i, resp.hosts[i].hostAddr.port);
+            ASSERT_EQ("offline", resp.hosts[i].status);
+        }
+    }
+}
+
 TEST(ProcessorTest, CreateSpaceTest) {
     fs::TempDir rootPath("/tmp/CreateSpaceTest.XXXXXX");
     std::unique_ptr<kvstore::KVStore> kv(TestUtils::initKV(rootPath.path()));