[Smartswitch] Fix incorrect reporting of data plane and control plane by DPU

sonic-chassisd/scripts/chassisd

@@ -776,12 +776,17 @@ class SmartSwitchModuleUpdater(ModuleUpdater):
     def update_dpu_state(self, key, state):
         """
         Update specific DPU state fields in chassisStateDB using the given key.
+        If state is 'down', delete the table first before setting new values.
         """
         try:
             # Connect to the CHASSIS_STATE_DB using daemon_base
             if not self.chassis_state_db:
                 self.chassis_state_db = daemon_base.db_connect("CHASSIS_STATE_DB")
 
+            # If state is down, delete the table first
+            if state == "down":
+                self.chassis_state_db.delete(key)
+
             # Prepare the fields to update
             updates = {
                 "dpu_midplane_link_state": state,
@@ -1367,12 +1372,14 @@ class DpuStateManagerTask(ProcessTaskBase):
         self.dpu_state_updater = dpu_state_updater
         self.state_db = daemon_base.db_connect('STATE_DB')
         self.app_db = daemon_base.db_connect('APPL_DB')
+        self.chassis_state_db = daemon_base.db_connect('CHASSIS_STATE_DB')
 
     def task_worker(self):
         sel = swsscommon.Select()
         selectable = [
             swsscommon.SubscriberStateTable(self.app_db, 'PORT_TABLE'),
-            swsscommon.SubscriberStateTable(self.state_db, 'SYSTEM_READY')
+            swsscommon.SubscriberStateTable(self.state_db, 'SYSTEM_READY'),
+            swsscommon.SubscriberStateTable(self.chassis_state_db, 'DPU_STATE')
         ]
 
         for s in selectable:
@@ -1388,10 +1395,31 @@ class DpuStateManagerTask(ProcessTaskBase):
                 if state != swsscommon.Select.OBJECT:
                     continue
 
-                for s in selectable:
-                    s.pops()
+                update_required = False
 
-                self.dpu_state_updater.update_state()
+                for s in selectable:
+                    result = s.pop()
+                    update_required = True # If there is any selectable object, we need to update the state
+                    if result is None:
+                        print("Here with None for {}".format(type(s)))
+                        continue
+                    key, op, fvp = result  # Changed from _ to fvp to match what we use below
+                    # Check if this is the DPU_STATE table
+                    if s.getDbConnector().getDbName() == 'CHASSIS_STATE_DB':
+                        # Don't update if this is a change for another DPU
+                        if key != self.dpu_state_updater.name:
+                            update_required = False
+                            continue
+                        # Don't update if this is our own dataplane/control plane state update to avoid recursion
+                        if op == 'SET' and isinstance(fvp, tuple):
+                            fvs = dict(fvp)
+                            if 'dpu_data_plane_state' in fvs and 'dpu_control_plane_state' in fvs:
+                                update_required = False
+                                continue
+                        self.logger.log_info(f"DPU_STATE change detected: operation={op}, key={key}")
+
+                if update_required:
+                    self.dpu_state_updater.update_state()
 
         except KeyboardInterrupt:
             pass

sonic-chassisd/tests/mock_swsscommon.py

@@ -74,6 +74,15 @@ def pop(self):
     def pops(self):
         return None
 
+    def getDbConnector(self):
+        return MockDbConnector()
+
+
+class MockDbConnector:
+
+    def getDbName(self):
+        return 'CHASSIS_STATE_DB'
+
 class RedisPipeline:
     def __init__(self, db):
         self.db = db

sonic-chassisd/tests/test_dpu_chassisd.py

@@ -233,3 +233,100 @@ def is_valid_date(date_str):
         daemon_chassisd.stop.wait.return_value = True
 
         thread.join()
+
+
+def test_dpu_state_manager_table_deletion():
+    chassis = MockDpuChassis()
+    chassis.get_dpu_id = MagicMock(return_value=0)
+    chassis.get_dataplane_state = MagicMock(return_value=True)
+    chassis.get_controlplane_state = MagicMock(return_value=True)
+
+    chassis_state_db = {}
+
+    def hset(key, field, value):
+        if key not in chassis_state_db:
+            chassis_state_db[key] = {}
+        chassis_state_db[key][field] = value
+
+    with mock.patch.object(swsscommon.Table, 'hset', side_effect=hset):
+        with mock.patch.object(swsscommon.SubscriberStateTable, 'pop', return_value=('DPU0', 'DEL', None)):
+            with mock.patch.object(swsscommon.Select, 'select',
+                side_effect=[(swsscommon.Select.OBJECT, None), KeyboardInterrupt]):
+
+                dpu_updater = DpuStateUpdater(SYSLOG_IDENTIFIER, chassis)
+                dpu_updater._time_now = MagicMock(return_value='Sat Jan 01 12:00:00 AM UTC 2000')
+
+                dpu_state_mng = DpuStateManagerTask(SYSLOG_IDENTIFIER, dpu_updater)
+                dpu_state_mng.task_worker()
+
+                # Verify state was updated since this is a DEL operation
+                assert chassis_state_db == {'DPU0': {
+                    'dpu_data_plane_state': 'up',
+                    'dpu_data_plane_time': 'Sat Jan 01 12:00:00 AM UTC 2000',
+                    'dpu_control_plane_state': 'up',
+                    'dpu_control_plane_time': 'Sat Jan 01 12:00:00 AM UTC 2000'
+                }}
+
+
+def test_dpu_state_manager_specific_key_update():
+    chassis = MockDpuChassis()
+    chassis.get_dpu_id = MagicMock(return_value=0)
+    chassis.get_dataplane_state = MagicMock(return_value=True)
+    chassis.get_controlplane_state = MagicMock(return_value=True)
+
+    chassis_state_db = {}
+
+    def hset(key, field, value):
+        if key not in chassis_state_db:
+            chassis_state_db[key] = {}
+        chassis_state_db[key][field] = value
+
+    with mock.patch.object(swsscommon.Table, 'hset', side_effect=hset):
+        with mock.patch.object(swsscommon.SubscriberStateTable, 'pop', return_value=('DPU0', 'SET', (('dpu_control_plane_state', 'up'), ('dpu_data_plane_state', 'up')))):
+            with mock.patch.object(swsscommon.Select, 'select',
+                side_effect=[(swsscommon.Select.OBJECT, None), KeyboardInterrupt]):
+
+                dpu_updater = DpuStateUpdater(SYSLOG_IDENTIFIER, chassis)
+                dpu_updater._time_now = MagicMock(return_value='Sat Jan 01 12:00:00 AM UTC 2000')
+
+                dpu_state_mng = DpuStateManagerTask(SYSLOG_IDENTIFIER, dpu_updater)
+                dpu_state_mng.task_worker()
+
+                # Verify no state update occurred since this was our own state update
+                assert chassis_state_db == {}
+
+
+def test_dpu_state_manager_none_result():
+    chassis = MockDpuChassis()
+    chassis.get_dpu_id = MagicMock(return_value=0)
+    chassis.get_dataplane_state = MagicMock(return_value=True)
+    chassis.get_controlplane_state = MagicMock(return_value=True)
+
+    chassis_state_db = {}
+
+    def hset(key, field, value):
+        if key not in chassis_state_db:
+            chassis_state_db[key] = {}
+        chassis_state_db[key][field] = value
+
+    def mock_pop():
+        return None
+
+    with mock.patch.object(swsscommon.Table, 'hset', side_effect=hset):
+        with mock.patch.object(swsscommon.SubscriberStateTable, 'pop', side_effect=mock_pop):
+            with mock.patch.object(swsscommon.Select, 'select',
+                side_effect=[(swsscommon.Select.OBJECT, None), KeyboardInterrupt]):
+
+                dpu_updater = DpuStateUpdater(SYSLOG_IDENTIFIER, chassis)
+                dpu_updater._time_now = MagicMock(return_value='Sat Jan 01 12:00:00 AM UTC 2000')
+
+                dpu_state_mng = DpuStateManagerTask(SYSLOG_IDENTIFIER, dpu_updater)
+                dpu_state_mng.task_worker()
+
+                # Verify state was updated even with None result
+                assert chassis_state_db == {'DPU0': {
+                    'dpu_data_plane_state': 'up',
+                    'dpu_data_plane_time': 'Sat Jan 01 12:00:00 AM UTC 2000',
+                    'dpu_control_plane_state': 'up',
+                    'dpu_control_plane_time': 'Sat Jan 01 12:00:00 AM UTC 2000'
+                }}