[Fast/Warm restart] Implement helper class for waiting restart done

common/Makefile.am

@@ -70,7 +70,8 @@ libswsscommon_la_SOURCES = \
     warm_restart.cpp          \
     luatable.cpp              \
     countertable.cpp          \
-    redisutility.cpp
+    redisutility.cpp          \
+    restart_waiter.cpp
 
 libswsscommon_la_CXXFLAGS = $(DBGFLAGS) $(AM_CFLAGS) $(CFLAGS_COMMON) $(LIBNL_CFLAGS) $(CODE_COVERAGE_CXXFLAGS)
 libswsscommon_la_CPPFLAGS = $(DBGFLAGS) $(AM_CFLAGS) $(CFLAGS_COMMON) $(LIBNL_CPPFLAGS) $(CODE_COVERAGE_CPPFLAGS)

common/restart_waiter.cpp

@@ -0,0 +1,130 @@
+#include "restart_waiter.h"
+#include "redispipeline.h"
+#include "select.h"
+#include "schema.h"
+#include "subscriberstatetable.h"
+#include "table.h"
+#include <string>
+
+using namespace swss;
+
+static const std::string STATE_DB_NAME = "STATE_DB";
+static const std::string STATE_DB_SEPARATOR = "|";
+static const std::string RESTART_KEY = "system";
+static const std::string RESTART_ENABLE_FIELD = "enable";
+static const std::string FAST_REBOOT_TABLE_NAME = "FAST_REBOOT";
+
+bool RestartWaiter::waitRestartDone(
+    unsigned int maxWaitSec,
+    unsigned int dbTimeout,
+    bool isTcpConn)
+{
+    DBConnector stateDb(STATE_DB_NAME, dbTimeout, isTcpConn);
+    return isWarmOrFastRestartInProgress(stateDb) ? doWait(stateDb, maxWaitSec) : true;
+}
+
+bool RestartWaiter::waitWarmRestartDone(unsigned int maxWaitSec,
+                                        unsigned int dbTimeout,
+                                        bool isTcpConn)
+{
+    DBConnector stateDb(STATE_DB_NAME, dbTimeout, isTcpConn);
+    if (isFastRestartInProgress(stateDb))
+    {
+        // It is fast boot, just return
+        return true;
+    }
+
+    return isWarmOrFastRestartInProgress(stateDb) ? doWait(stateDb, maxWaitSec) : true;
+}
+
+bool RestartWaiter::waitFastRestartDone(unsigned int maxWaitSec,
+                                        unsigned int dbTimeout,
+                                        bool isTcpConn)
+{
+    DBConnector stateDb(STATE_DB_NAME, dbTimeout, isTcpConn);
+    if (!isFastRestartInProgress(stateDb))
+    {
+        // Fast boot is not in progress
+        return true;
+    }
+
+    return isWarmOrFastRestartInProgress(stateDb) ? doWait(stateDb, maxWaitSec) : true;
+}
+
+bool RestartWaiter::doWait(DBConnector &stateDb,
+                           unsigned int maxWaitSec)
+{
+    int selectTimeout = static_cast<int>(maxWaitSec);
+
+    SubscriberStateTable restartEnableTable(&stateDb, STATE_WARM_RESTART_ENABLE_TABLE_NAME);
+    Select s;
+    s.addSelectable(&restartEnableTable);
+
+    auto start = std::chrono::steady_clock::now();
+    while (1)
+    {
+        Selectable *sel = NULL;
+        int ret = s.select(&sel, selectTimeout * 1000, true);
+
+        if (ret == Select::OBJECT)
+        {
+            KeyOpFieldsValuesTuple kco;
+            restartEnableTable.pop(kco);
+            auto &key = kfvKey(kco);
+            if (key == RESTART_KEY)
+            {
+                auto& values = kfvFieldsValues(kco);
+                for (auto& fvt: values)
+                {
+                    auto& field = fvField(fvt);
+                    auto& value = fvValue(fvt);
+                    if (field == RESTART_ENABLE_FIELD)
+                    {
+                        if (value == "false")
+                        {
+                            return true;
+                        }
+                        else
+                        {
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+        else if (ret == Select::ERROR)
+        {
+            SWSS_LOG_NOTICE("Error: wait restart done got error - %s!", strerror(errno));
+        }
+        else if (ret == Select::TIMEOUT)
+        {
+            SWSS_LOG_INFO("Timeout: wait restart done got select timeout");
+        }
+        else if (ret == Select::SIGNALINT)
+        {
+            return false;
+        }
+
+        auto end = std::chrono::steady_clock::now();
+        int delay = static_cast<int>(
+                    std::chrono::duration_cast<std::chrono::seconds>(end - start).count());
+
+        selectTimeout -= delay;
+        if (selectTimeout < 0)
+        {
+            return false;
+        }
+    }
+}
+
+bool RestartWaiter::isWarmOrFastRestartInProgress(DBConnector &stateDb)
+{
+    auto ret = stateDb.hget(STATE_WARM_RESTART_ENABLE_TABLE_NAME + STATE_DB_SEPARATOR + RESTART_KEY, RESTART_ENABLE_FIELD);
+    return ret && *ret.get() == "true";
+}
+
+bool RestartWaiter::isFastRestartInProgress(DBConnector &stateDb)
+{
+    auto ret = stateDb.get(FAST_REBOOT_TABLE_NAME + STATE_DB_SEPARATOR + RESTART_KEY);
+    return ret.get() != nullptr;
+}

common/restart_waiter.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "dbconnector.h"
+
+namespace swss
+{
+
+// Helper class to wait for warm/fast reboot done
+class RestartWaiter
+{
+public:
+    static bool waitRestartDone(unsigned int maxWaitSec = 180,
+                                unsigned int dbTimeout = 0,
+                                bool isTcpConn = false);
+
+    static bool waitWarmRestartDone(unsigned int maxWaitSec = 180,
+                                    unsigned int dbTimeout = 0,
+                                    bool isTcpConn = false);
+
+    static bool waitFastRestartDone(unsigned int maxWaitSec = 180,
+                                    unsigned int dbTimeout = 0,
+                                    bool isTcpConn = false);
+
+private:
+    static bool doWait(swss::DBConnector &stateDb,
+                       unsigned int maxWaitSec);
+
+    static bool isWarmOrFastRestartInProgress(swss::DBConnector &stateDb);
+    static bool isFastRestartInProgress(swss::DBConnector &stateDb);
+};
+
+}

pyext/swsscommon.i

@@ -38,6 +38,7 @@
 #include "events.h"
 #include "configdb.h"
 #include "status_code_util.h"
+#include "restart_waiter.h"
 %}
 
 %include <std_string.i>
@@ -219,3 +220,4 @@ T castSelectableObj(swss::Selectable *temp)
 %include "events.h"
 
 %include "status_code_util.h"
+%include "restart_waiter.h"

tests/Makefile.am

@@ -41,6 +41,7 @@ tests_SOURCES = redis_ut.cpp                \
                 events_common_ut.cpp        \
                 events_service_ut.cpp       \
                 events_ut.cpp               \
+                restart_waiter_ut.cpp       \
                 main.cpp
 
 tests_CFLAGS = $(DBGFLAGS) $(AM_CFLAGS) $(CFLAGS_COMMON) $(CFLAGS_GTEST) $(LIBNL_CFLAGS)

tests/restart_waiter_ut.cpp

@@ -0,0 +1,113 @@
+#include <gtest/gtest.h>
+#include <thread>
+#include <string>
+#include <unistd.h>
+#include <vector>
+
+#include "common/dbconnector.h"
+#include "common/restart_waiter.h"
+#include "common/schema.h"
+#include "common/table.h"
+
+using namespace swss;
+using namespace std;
+
+static const string FAST_REBOOT_KEY = "FAST_REBOOT|system";
+
+static void set_reboot_status(string status, int delay = 0)
+{
+    if (delay > 0)
+    {
+        sleep(delay);
+    }
+
+    DBConnector db("STATE_DB", 0);
+    Table table(&db, STATE_WARM_RESTART_ENABLE_TABLE_NAME);
+    table.hset("system", "enable", status);
+}
+
+class FastBootHelper
+{
+public:
+    FastBootHelper(): db("STATE_DB", 0)
+    {
+        db.set(FAST_REBOOT_KEY, "1");
+    }
+
+    ~FastBootHelper()
+    {
+        db.del({FAST_REBOOT_KEY});
+    }
+private:
+    DBConnector db;
+};
+
+TEST(RestartWaiter, success)
+{
+    set_reboot_status("true");
+    thread t(set_reboot_status, "false", 3);
+    EXPECT_TRUE(RestartWaiter::waitRestartDone());
+    t.join();
+}
+
+TEST(RestartWaiter, successWarmRestart)
+{
+    set_reboot_status("true");
+    thread t(set_reboot_status, "false", 3);
+    EXPECT_TRUE(RestartWaiter::waitWarmRestartDone());
+    t.join();
+}
+
+TEST(RestartWaiter, successFastRestart)
+{
+    FastBootHelper helper;
+    set_reboot_status("true");
+    thread t(set_reboot_status, "false", 3);
+    EXPECT_TRUE(RestartWaiter::waitFastRestartDone());
+    t.join();
+}
+
+TEST(RestartWaiter, timeout)
+{
+    set_reboot_status("true");
+    EXPECT_FALSE(RestartWaiter::waitRestartDone(1));
+    EXPECT_FALSE(RestartWaiter::waitWarmRestartDone(1));
+
+    FastBootHelper helper;
+    EXPECT_FALSE(RestartWaiter::waitFastRestartDone(1));
+
+    set_reboot_status("false");
+}
+
+TEST(RestartWaiter, successNoDelay)
+{
+    set_reboot_status("false");
+    EXPECT_TRUE(RestartWaiter::waitRestartDone());
+    EXPECT_TRUE(RestartWaiter::waitWarmRestartDone());
+
+    FastBootHelper helper;
+    EXPECT_TRUE(RestartWaiter::waitFastRestartDone());
+}
+
+TEST(RestartWaiter, successNoKey)
+{
+    DBConnector db("STATE_DB", 0);
+    string key = string(STATE_WARM_RESTART_ENABLE_TABLE_NAME) + string("|system");
+    db.del({key});
+    EXPECT_TRUE(RestartWaiter::waitRestartDone());
+    EXPECT_TRUE(RestartWaiter::waitWarmRestartDone());
+
+    FastBootHelper helper;
+    EXPECT_TRUE(RestartWaiter::waitFastRestartDone());
+}
+
+TEST(RestartWaiter, waitWarmButFastInProgress)
+{
+    FastBootHelper helper;
+    EXPECT_TRUE(RestartWaiter::waitWarmRestartDone());
+}
+
+TEST(RestartWaiter, waitFastButFastNotInProgress)
+{
+    EXPECT_TRUE(RestartWaiter::waitFastRestartDone());
+}