Improve reboot status tracking and add reboot count & method to status (#259)

Why I did it
To provide more accurate and detailed tracking of reboot operations on the host. This includes:

Fixing incorrect or missing values in the reboot_status_flag.
Adding the ability to track how many times a reboot has been initiated (count).
Recording which reboot method was used (method), which is critical for debugging and telemetry.
Fixes: sonic-net/sonic-buildimage#22157

How I did it
Enhanced the populate_reboot_status_flag() function to also track:
-- count: Number of reboots initiated by the host service.
-- method: Type of reboot (e.g., cold, warm, halt).

Corrected how container status is evaluated using Docker SDK by checking attrs['State']['Running'] instead of relying on the status string.

Updated all relevant call sites of populate_reboot_status_flag() to pass the correct parameters.

Adjusted unit tests to validate updated logic and ensure coverage of new parameters.

Author: vvolam

host_modules/reboot.py

@@ -8,6 +8,7 @@
 import psutil
 from host_modules import host_service
 from utils.run_cmd import _run_command
+from enum import Enum
 
 MOD_NAME = 'reboot'
 # Reboot method in reboot request
@@ -25,6 +26,12 @@
 EXECUTE_HALT_REBOOT_COMMAND = "sudo reboot -p"
 EXECUTE_WARM_REBOOT_COMMAND = "sudo warm-reboot"
 
+class RebootStatus(Enum):
+    STATUS_UNKNOWN = 0
+    STATUS_SUCCESS = 1
+    STATUS_RETRIABLE_FAILURE = 2
+    STATUS_FAILURE = 3
+
 logger = logging.getLogger(__name__)
 
 
@@ -38,16 +45,23 @@ def __init__(self, mod_name):
         self.lock = threading.Lock()
         # reboot_status_flag is used to keep track of reboot status on host
         self.reboot_status_flag = {}
+
+        # reboot count
+        self.reboot_count = 0
+
         # Populating with default value i.e., no active reboot
         self.populate_reboot_status_flag()
         super(Reboot, self).__init__(mod_name)
 
-    def populate_reboot_status_flag(self, active = False, when = 0, reason = ""):
+    def populate_reboot_status_flag(self, active = False, when = 0, reason = "", method = "", status = RebootStatus.STATUS_UNKNOWN):
         """Populates the reboot_status_flag with given input params"""
         self.lock.acquire()
         self.reboot_status_flag["active"] = active
         self.reboot_status_flag["when"] = when
         self.reboot_status_flag["reason"] = reason
+        self.reboot_status_flag["count"] = self.reboot_count
+        self.reboot_status_flag["method"] = method
+        self.reboot_status_flag["status"] = status
         self.lock.release()
         return
 
@@ -71,10 +85,10 @@ def is_container_running(self, container_name):
         """Check if a given container is running using the Docker SDK."""
         try:
             client = docker.from_env()
-            containers = client.containers.list(filters={"name": container_name})
+            containers = client.containers.list(filters={"name": f"^{container_name}$"})
 
             for container in containers:
-                if container.name == container_name and container.status == "running":
+                if container.name == container_name and container.attrs['State']['Running']:
                     return True
             return False
         except Exception as e:
@@ -111,7 +125,7 @@ def execute_reboot(self, reboot_method):
 
         rc, stdout, stderr = _run_command(command)
         if rc:
-            self.populate_reboot_status_flag()
+            self.populate_reboot_status_flag(False, int (time.time()), "Failed to execute reboot command", reboot_method, RebootStatus.STATUS_FAILURE)
             logger.error("%s: Reboot failed execution with stdout: %s, "
                          "stderr: %s", MOD_NAME, stdout, stderr)
             return
@@ -128,29 +142,30 @@ def execute_reboot(self, reboot_method):
             logger.info("%s: Waiting until services are halted or timeout occurs", MOD_NAME)
             timeout = HALT_TIMEOUT
             start_time = time.monotonic()
+
             while time.monotonic() - start_time < timeout:
                 if not self.is_halt_command_running() and not self.is_container_running("pmon"):
                     logger.info("%s: Halting the services is completed on the device", MOD_NAME)
+                    self.populate_reboot_status_flag(False, 0, "Halt reboot completed", reboot_method, RebootStatus.STATUS_SUCCESS)
                     return
                 time.sleep(5)
 
             # Check if PMON container is still running after timeout
             if self.is_halt_command_running() or self.is_container_running("pmon"):
                 #Halt reboot has failed, as pmon is still running.
+                self.populate_reboot_status_flag(False, int(time.time()), "Halt reboot did not complete", reboot_method, RebootStatus.STATUS_FAILURE)
                 logger.error("%s: HALT reboot failed: Services are still running", MOD_NAME)
-                self.populate_reboot_status_flag()
-                return
             else:
+                self.populate_reboot_status_flag(False, 0, "Halt reboot completed", reboot_method, RebootStatus.STATUS_SUCCESS)
                 logger.info("%s: Halting the services is completed on the device", MOD_NAME)
-
+            return
         else:
             time.sleep(REBOOT_TIMEOUT)
             # Conclude that the reboot has failed if we reach this point
-            self.populate_reboot_status_flag()
+            self.populate_reboot_status_flag(False, int(time.time()), "Reboot command failed to execute", reboot_method, RebootStatus.STATUS_FAILURE)
             return
 
     @host_service.method(host_service.bus_name(MOD_NAME), in_signature='as', out_signature='is')
-
     def issue_reboot(self, options):
         """Initializes reboot thorugh RPC based on the reboot flag assigned.
            Issues reboot after performing the following steps sequentially:
@@ -162,6 +177,7 @@ def issue_reboot(self, options):
         logger.warning("%s: issue_reboot rpc called", MOD_NAME)
         self.lock.acquire()
         is_reboot_ongoing = self.reboot_status_flag["active"]
+        self.reboot_count += 1
         self.lock.release()
         # Return without issuing the reboot if the previous reboot is ongoing
         if is_reboot_ongoing:
@@ -185,7 +201,7 @@ def issue_reboot(self, options):
             return err, errstr
 
         # Sets reboot_status_flag to be in active state
-        self.populate_reboot_status_flag(True, int(time.time()), reboot_request.get("message", ""))
+        self.populate_reboot_status_flag(True, int(time.time()), reboot_request.get("message", ""), reboot_request["method"], RebootStatus.STATUS_UNKNOWN)
 
         # Issue reboot in a new thread and reset the reboot_status_flag if the reboot fails
         try:

tests/host_modules/reboot_test.py

@@ -4,6 +4,7 @@
 import json
 import sys
 import os
+import time
 import pytest
 import logging
 
@@ -19,6 +20,8 @@
 host_modules_path = os.path.join(sonic_host_service_path, "../host_modules")
 sys.path.insert(0, sonic_host_service_path)
 
+from host_modules.reboot import RebootStatus
+
 TIME = 1617811205
 TEST_ACTIVE_RESPONSE_DATA = "{\"active\": true, \"when\": 1617811205, \"reason\": \"testing reboot response\"}"
 TEST_INACTIVE_RESPONSE_DATA = "{\"active\": false, \"when\": 0, \"reason\": \"\"}"
@@ -54,6 +57,9 @@ def test_populate_reboot_status_flag(self):
             assert self.reboot_module.reboot_status_flag["active"] == False
             assert self.reboot_module.reboot_status_flag["when"] == 0
             assert self.reboot_module.reboot_status_flag["reason"] == ""
+            assert self.reboot_module.reboot_status_flag["count"] == 0
+            assert self.reboot_module.reboot_status_flag["method"] == ""
+            assert self.reboot_module.reboot_status_flag["status"] == RebootStatus.STATUS_UNKNOWN
 
     def test_validate_reboot_request_success_cold_boot_enum_method(self):
         reboot_request = {"method": REBOOT_METHOD_COLD_BOOT_ENUM, "reason": "test reboot request reason"}
@@ -115,12 +121,12 @@ def test_is_container_running_success(self):
             mock_docker.return_value = mock_client
             mock_container = mock.Mock()
             mock_container.name = "pmon"
-            mock_container.status = "running"
+            mock_container.attrs = {'State': {'Running': True}}
             mock_client.containers.list.return_value = [mock_container]
 
             result = self.reboot_module.is_container_running("pmon")
             assert result is True
-            mock_client.containers.list.assert_called_once_with(filters={"name": "pmon"})
+            mock_client.containers.list.assert_called_once_with(filters={"name": "^pmon$"})
 
     def test_is_container_running_failure(self):
         with mock.patch("docker.from_env") as mock_docker:
@@ -130,7 +136,7 @@ def test_is_container_running_failure(self):
 
             result = self.reboot_module.is_container_running("pmon")
             assert result is False
-            mock_client.containers.list.assert_called_once_with(filters={"name": "pmon"})
+            mock_client.containers.list.assert_called_once_with(filters={"name": "^pmon$"})
 
     def test_is_container_running_exception(self, caplog):
         with mock.patch("docker.from_env", side_effect=Exception("Docker error")) as mock_docker, \
@@ -176,7 +182,7 @@ def test_execute_reboot_success(self):
             self.reboot_module.execute_reboot("WARM")
             mock_run_command.assert_called_once_with("sudo warm-reboot")
             mock_sleep.assert_called_once_with(260)
-            mock_populate_reboot_status_flag.assert_called_once_with()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, int(time.time()), "Reboot command failed to execute", 'WARM', RebootStatus.STATUS_FAILURE)
 
     def test_execute_reboot_fail_unknown_reboot(self, caplog):
         with caplog.at_level(logging.ERROR):
@@ -196,7 +202,7 @@ def test_execute_reboot_fail_issue_reboot_command_cold_boot(self, caplog):
                     "stdout: ['stdout: execute cold reboot'], stderr: "
                     "['stderror: execute cold reboot']")
             assert caplog.records[0].message == msg
-            mock_populate_reboot_status_flag.assert_called_once_with()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, int (time.time()), "Failed to execute reboot command", 1, RebootStatus.STATUS_FAILURE)
 
     def test_execute_reboot_fail_issue_reboot_command_halt(self, caplog):
         with (
@@ -210,7 +216,7 @@ def test_execute_reboot_fail_issue_reboot_command_halt(self, caplog):
                    "stdout: ['stdout: execute halt reboot'], stderr: "
                    "['stderror: execute halt reboot']")
             assert caplog.records[0].message == msg
-            mock_populate_reboot_status_flag.assert_called_once_with()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, int (time.time()), "Failed to execute reboot command", 3, RebootStatus.STATUS_FAILURE)
 
     def test_execute_reboot_success_halt(self):
         with (
@@ -225,7 +231,7 @@ def test_execute_reboot_success_halt(self):
             mock_run_command.assert_called_once_with("sudo reboot -p")
             mock_is_halt_command_running.assert_called()
             mock_is_container_running.assert_called_with("pmon")
-            mock_populate_reboot_status_flag.assert_not_called()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, 0, 'Halt reboot completed', 3, RebootStatus.STATUS_SUCCESS)
 
     def test_execute_reboot_fail_halt_timeout(self, caplog):
         with (
@@ -242,7 +248,7 @@ def test_execute_reboot_fail_halt_timeout(self, caplog):
             mock_sleep.assert_called_with(5)
             mock_is_halt_command_running.assert_called()
             assert any("HALT reboot failed: Services are still running" in record.message for record in caplog.records)
-            mock_populate_reboot_status_flag.assert_called_once_with()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, int(time.time()), 'Halt reboot did not complete', 3, RebootStatus.STATUS_FAILURE)
 
     def test_execute_reboot_fail_issue_reboot_command_warm(self, caplog):
         with (
@@ -256,7 +262,7 @@ def test_execute_reboot_fail_issue_reboot_command_warm(self, caplog):
                     "stdout: ['stdout: execute WARM reboot'], stderr: "
                     "['stderror: execute WARM reboot']")
             assert caplog.records[0].message == msg
-            mock_populate_reboot_status_flag.assert_called_once_with()
+            mock_populate_reboot_status_flag.assert_called_once_with(False, int (time.time()), "Failed to execute reboot command", 'WARM', RebootStatus.STATUS_FAILURE)
 
     def test_issue_reboot_success_cold_boot(self):
         with (
@@ -355,22 +361,26 @@ def test_issue_reboot_fail_issue_reboot_thread(self):
 
     def test_get_reboot_status_active(self):
         MSG="testing reboot response"
-        self.reboot_module.populate_reboot_status_flag(True, TIME, MSG)
+        self.reboot_module.populate_reboot_status_flag(True, TIME, MSG, REBOOT_METHOD_COLD_BOOT_ENUM, RebootStatus.STATUS_SUCCESS.name)
         result = self.reboot_module.get_reboot_status()
         assert result[0] == 0
         response_data = json.loads(result[1])
         assert response_data["active"] == True
         assert response_data["when"] == TIME
         assert response_data["reason"] == MSG
+        assert response_data["method"] == REBOOT_METHOD_COLD_BOOT_ENUM
+        assert response_data["status"] == RebootStatus.STATUS_SUCCESS.name
 
     def test_get_reboot_status_inactive(self):
-        self.reboot_module.populate_reboot_status_flag(False, 0, "")
+        self.reboot_module.populate_reboot_status_flag(False, 0, "", REBOOT_METHOD_COLD_BOOT_ENUM, RebootStatus.STATUS_SUCCESS.name)
         result = self.reboot_module.get_reboot_status()
         assert result[0] == 0
         response_data = json.loads(result[1])
         assert response_data["active"] == False
         assert response_data["when"] == 0
         assert response_data["reason"] == ""
+        assert response_data["method"] == REBOOT_METHOD_COLD_BOOT_ENUM
+        assert response_data["status"] == RebootStatus.STATUS_SUCCESS.name
 
 #        assert result[1] == TEST_INACTIVE_RESPONSE_DATA