From 0b2d62e38c301be23bb07234032d7610a543c896 Mon Sep 17 00:00:00 2001
From: PFZheng <wenffie@gmail.com>
Date: Tue, 18 Feb 2020 19:52:27 -0800
Subject: [PATCH] fix LeaseTest::apply_thread_hung unit test

---
 src/braft/test/util.h      |  6 +++++-
 test/test_leader_lease.cpp | 21 ++++++++++++---------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/braft/test/util.h b/src/braft/test/util.h
index 9856045c..210c241e 100644
--- a/src/braft/test/util.h
+++ b/src/braft/test/util.h
@@ -230,7 +230,8 @@ class Cluster {
     }
 
     int start(const butil::EndPoint& listen_addr, bool empty_peers = false,
-              int snapshot_interval_s = 30) {
+              int snapshot_interval_s = 30,
+              braft::Closure* leader_start_closure = NULL) {
         if (_server_map[listen_addr] == NULL) {
             brpc::Server* server = new brpc::Server();
             if (braft::add_service(server, listen_addr) != 0 
@@ -250,6 +251,9 @@ class Cluster {
             options.initial_conf = braft::Configuration(_peers);
         }
         MockFSM* fsm = new MockFSM(listen_addr);
+        if (leader_start_closure) {
+            fsm->set_on_leader_start_closure(leader_start_closure);
+        }
         options.fsm = fsm;
         options.node_owns_fsm = true;
         butil::string_printf(&options.log_uri, "local://./data/%s/log",
diff --git a/test/test_leader_lease.cpp b/test/test_leader_lease.cpp
index 565af854..723ae22c 100644
--- a/test/test_leader_lease.cpp
+++ b/test/test_leader_lease.cpp
@@ -477,12 +477,12 @@ class OnLeaderStartHungClosure : public braft::SynchronizedClosure {
 
     void Run() {
         while (hung) {
-            usleep(10 * 1000);
+            bthread_usleep(10 * 1000);
         }
         braft::SynchronizedClosure::Run();
     }
 
-   volatile bool hung;
+    butil::atomic<bool> hung;
 };
 
 TEST_F(LeaseTest, apply_thread_hung) {
@@ -502,21 +502,23 @@ TEST_F(LeaseTest, apply_thread_hung) {
     // start cluster
     Cluster cluster("unittest", peers, 500, 10);
     for (size_t i = 0; i < peers.size(); i++) {
-        ASSERT_EQ(0, cluster.start(peers[i].addr));
+        ASSERT_EQ(0, cluster.start(peers[i].addr, false, 30, on_leader_start_closures[i]));
     }
 
-    std::vector<braft::Node*> nodes;
+    /*
     cluster.all_nodes(&nodes);
     for (size_t i = 0; i < nodes.size(); ++i) {
         static_cast<MockFSM*>(nodes[i]->_impl->_options.fsm)
             ->set_on_leader_start_closure(on_leader_start_closures[i]);
     }
+    */
 
     // elect leader
     cluster.wait_leader();
     braft::Node* leader = cluster.leader();
     ASSERT_TRUE(leader != NULL);
     LOG(WARNING) << "leader is " << leader->node_id();
+    std::vector<braft::Node*> nodes;
     cluster.followers(&nodes);
     braft::PeerId target = nodes[0]->node_id().peer_id;
 
@@ -534,11 +536,12 @@ TEST_F(LeaseTest, apply_thread_hung) {
         on_leader_start_closures[i]->hung = false;
     }
 
-    usleep(100 * 1000);
-
-    // lease status should be valid
-    leader->get_leader_lease_status(&lease_status);
-    ASSERT_EQ(lease_status.state, braft::LEASE_VALID);
+    for (int i = 0; i < 10; ++i) {
+        usleep(100 * 1000);
+        // lease status should be valid
+        leader->get_leader_lease_status(&lease_status);
+        ASSERT_EQ(lease_status.state, braft::LEASE_VALID);
+    }
 
     cluster.stop_all();