[SmartSwitch] Add graceful shutdown and startup handling in platform daemons

sonic-chassisd/scripts/chassisd

@@ -254,16 +254,7 @@ class SmartSwitchModuleConfigUpdater(logger.Logger):
             self.log_warning("Invalid admin_state value: {}".format(admin_state))
 
     def submit_callback(self, module_index, admin_state, key):
-        if admin_state == MODULE_ADMIN_DOWN:
-            # This is only valid on platforms which have pci_detach and sensord changes required. If it is not implemented,
-            # there are no actions taken during this function execution.
-            try_get(self.chassis.get_module(module_index).module_pre_shutdown, default=False)
-        try_get(self.chassis.get_module(module_index).set_admin_state, admin_state, default=False)
-        if admin_state == MODULE_ADMIN_UP:
-            # This is only valid on platforms which have pci_rescan sensord changes required. If it is not implemented,
-            # there are no actions taken during this function execution.
-            try_get(self.chassis.get_module(module_index).module_post_startup, default=False)
-        pass
+        try_get(self.chassis.get_module(module_index).set_admin_state_gracefully, admin_state, default=False)
 
 #
 # Module Updater ==============================================================
@@ -723,7 +714,6 @@ class SmartSwitchModuleUpdater(ModuleUpdater):
         self.module_reboot_table = swsscommon.Table(self.chassis_state_db, CHASSIS_MODULE_REBOOT_INFO_TABLE)
         self.down_modules = {}
         self.chassis_app_db_clean_sha = None
-        self.module_transition_flag_helper = ModuleTransitionFlagHelper()
 
         self.midplane_initialized = try_get(chassis.init_midplane_switch, default=False)
         if not self.midplane_initialized:
@@ -815,9 +805,6 @@ class SmartSwitchModuleUpdater(ModuleUpdater):
                     # Persist dpu down time
                     self.persist_dpu_reboot_time(key)
                     # persist reboot cause
-                    # Clear transition flag in STATE_DB
-                    self.module_transition_flag_helper.clear_transition_flag(key)
-
                     reboot_cause = try_get(self.chassis.get_module(module_index).get_reboot_cause)
                     self.persist_dpu_reboot_cause(reboot_cause, key)
                     # publish reboot cause to db
@@ -852,9 +839,6 @@ class SmartSwitchModuleUpdater(ModuleUpdater):
                         self.persist_dpu_reboot_cause(reboot_cause, key)
                         self.update_dpu_reboot_cause_to_db(key)
 
-                    # Clear transition flag in STATE_DB
-                    self.module_transition_flag_helper.clear_transition_flag(key)
-
     def _get_module_info(self, module_index):
         """
         Retrieves module info of this module
@@ -1336,34 +1320,6 @@ class DpuStateUpdater(logger.Logger):
         self._update_dp_dpu_state('down')
         self._update_cp_dpu_state('down')
 
-class ModuleTransitionFlagHelper(logger.Logger):
-    def __init__(self, log_identifier = SYSLOG_IDENTIFIER):
-        super(ModuleTransitionFlagHelper, self).__init__(log_identifier)
-        # Use new connector to avoid redis failures
-        """Create a helper function to get the module table,
-        since multiple threads updating with the same connector will cause redis failures"""
-        state_db = daemon_base.db_connect("STATE_DB")
-        self.module_table = swsscommon.Table(state_db, CHASSIS_MODULE_INFO_TABLE)
-
-    def set_transition_flag(self, module_name):
-        try:
-            self.module_table.hset(module_name, 'state_transition_in_progress', 'True')
-            self.module_table.hset(module_name, 'transition_start_time', datetime.now(timezone.utc).replace(tzinfo=None).isoformat())
-        except Exception as e:
-            self.log_error(f"Error setting transition flag for {module_name}: {e}")
-
-    def clear_transition_flag(self, module_name):
-        try:
-            self.log_info(f"Clearing transition flag for {module_name}")
-            self.module_table.hdel(module_name, 'state_transition_in_progress')
-            self.module_table.hdel(module_name, 'transition_start_time')
-        except Exception as e:
-            self.log_error(f"Error clearing transition flag for {module_name}: {e}")
-
-    def clear_all_transition_flags(self):
-        for module_name in self.module_table.getKeys():
-            self.clear_transition_flag(module_name)
-
 #
 # Daemon =======================================================================
 #
@@ -1406,8 +1362,7 @@ class ChassisdDaemon(daemon_base.DaemonBase):
         try_get(self.module_updater.chassis.get_module(module_index).module_pre_shutdown, default=False)
         # Set admin_state change in progress using the centralized method
         if admin_state == MODULE_ADMIN_DOWN:
-            ModuleTransitionFlagHelper().set_transition_flag(module_name)
-            try_get(self.module_updater.chassis.get_module(module_index).set_admin_state, admin_state, default=False)
+            try_get(self.module_updater.chassis.get_module(module_index).set_admin_state_gracefully, admin_state, default=False)
 
     def set_initial_dpu_admin_state(self):
         """Send admin_state trigger once to modules those are powered up"""
@@ -1486,16 +1441,7 @@ class ChassisdDaemon(daemon_base.DaemonBase):
 
         # Set the initial DPU admin state for SmartSwitch
         if self.smartswitch:
-            # Clear all stale transition flags for SmartSwitch on startup
-            ModuleTransitionFlagHelper().clear_all_transition_flags()
-            self.set_initial_dpu_admin_state()
-            # Clear all transition flags for SmartSwitch after setting the initial DPU admin state
-            module_transition_flag_helper = ModuleTransitionFlagHelper()
-            # Clear all stale transition flags for SmartSwitch on startup
-            module_transition_flag_helper.clear_all_transition_flags()
             self.set_initial_dpu_admin_state()
-            # Clear all transition flags for SmartSwitch after setting the initial DPU admin state
-            module_transition_flag_helper.clear_all_transition_flags()
 
         while not self.stop.wait(CHASSIS_INFO_UPDATE_PERIOD_SECS):
             self.module_updater.module_db_update()

sonic-chassisd/tests/mock_platform.py

@@ -78,6 +78,10 @@ def module_pre_shutdown(self):
     def module_post_startup(self):
         pass
 
+    def set_admin_state_gracefully(self, up):
+        """Mock implementation of set_admin_state_gracefully"""
+        return self.set_admin_state(up)
+
     def is_midplane_reachable(self):
         return self.midplane_access
 

sonic-chassisd/tests/test_chassisd.py

@@ -628,19 +628,15 @@ def test_smartswitch_configupdater_check_admin_state():
 
     # Test setting admin state to down
     admin_state = 0
-    with patch.object(module, 'module_pre_shutdown') as mock_module_pre_shutdown, \
-         patch.object(module, 'set_admin_state') as mock_set_admin_state:
+    with patch.object(module, 'set_admin_state_gracefully') as mock_set_admin_state_gracefully:
         config_updater.module_config_update(name, admin_state)
-        mock_module_pre_shutdown.assert_called_once()
-        mock_set_admin_state.assert_called_once_with(admin_state)
+        mock_set_admin_state_gracefully.assert_called_once_with(admin_state)
 
     # Test setting admin state to up
     admin_state = 1
-    with patch.object(module, 'set_admin_state') as mock_set_admin_state, \
-         patch.object(module, 'module_post_startup') as mock_module_post_startup:
+    with patch.object(module, 'set_admin_state_gracefully') as mock_set_admin_state_gracefully:
         config_updater.module_config_update(name, admin_state)
-        mock_set_admin_state.assert_called_once_with(admin_state)
-        mock_module_post_startup.assert_called_once()
+        mock_set_admin_state_gracefully.assert_called_once_with(admin_state)
 
 
 @patch("chassisd.glob.glob")
@@ -1966,24 +1962,19 @@ def test_submit_dpu_callback():
 
     # Test MODULE_ADMIN_DOWN scenario
     with patch.object(module, 'module_pre_shutdown') as mock_pre_shutdown, \
-         patch.object(module, 'set_admin_state') as mock_set_admin_state, \
-         patch.object(module, 'module_post_startup') as mock_post_startup:
+         patch.object(module, 'set_admin_state_gracefully') as mock_set_admin_state_gracefully:
         daemon_chassisd.submit_dpu_callback(index, MODULE_ADMIN_DOWN, name)
         # Verify correct functions are called for admin down
         mock_pre_shutdown.assert_called_once()
-        mock_set_admin_state.assert_called_once_with(MODULE_ADMIN_DOWN)
-        mock_post_startup.assert_not_called()
+        mock_set_admin_state_gracefully.assert_called_once_with(MODULE_ADMIN_DOWN)
 
-
-    # Reset mocks for next test
+    # Test non-MODULE_ADMIN_DOWN scenario (e.g., MODULE_PRE_SHUTDOWN)
     with patch.object(module, 'module_pre_shutdown') as mock_pre_shutdown, \
-         patch.object(module, 'set_admin_state') as mock_set_admin_state, \
-         patch.object(module, 'module_post_startup') as mock_post_startup:
+         patch.object(module, 'set_admin_state_gracefully') as mock_set_admin_state_gracefully:
 
         module_updater.module_table.get = MagicMock(return_value=(True, []))
         daemon_chassisd.submit_dpu_callback(index, MODULE_PRE_SHUTDOWN, name)
 
-        # Verify correct functions are called for admin up
+        # Verify only pre_shutdown is called for non-admin-down states
         mock_pre_shutdown.assert_called_once()
-        mock_set_admin_state.assert_not_called()
-        mock_post_startup.assert_not_called()
+        mock_set_admin_state_gracefully.assert_not_called()