Skip to content

gnoi-shutdown: use per-thread DB connections to avoid crossed DPU reads #408

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gpunathilell wants to merge 2 commits into
base: master
Choose a base branch
from
+89 −16
Open

gnoi-shutdown: use per-thread DB connections to avoid crossed DPU reads #408

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aae4414
gnoi-shutdown: use per-thread DB connections to avoid crossed DPU reads
gpunathilell Jul 10, 2026
c6b49a9
gnoi-shutdown: remove unused transition_type param; add per-thread DB…
gpunathilell Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions scripts/gnoi_shutdown_daemon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,11 +142,18 @@ def _should_skip_gnoi_shutdown(self, dpu_name: str):
ModuleBase.MODULE_STATUS_POWERED_DOWN,
)

def _handle_transition(self, dpu_name: str, transition_type: str) -> bool:
def _handle_transition(self, dpu_name: str,
config_db=None, state_db=None) -> bool:
"""
Handle a shutdown or reboot transition for a DPU module.
Returns True if the operation completed successfully, False otherwise.

config_db/state_db are per-thread DB connections; fall back to the
shared connections for direct/unit-test callers.
"""
config_db = config_db if config_db is not None else self._config_db
state_db = state_db if state_db is not None else self._db

logger.log_notice(f"{dpu_name}: Starting gNOI shutdown sequence")

# Check if DPU is already powered off / offline before attempting gNOI shutdown.
Expand Down Expand Up @@ -174,14 +181,14 @@ def _handle_transition(self, dpu_name: str, transition_type: str) -> bool:
return cleared

# Wait for platform PCI detach completion
if not self._wait_for_gnoi_halt_in_progress(dpu_name):
if not self._wait_for_gnoi_halt_in_progress(dpu_name, state_db):
logger.log_warning(f"{dpu_name}: Timeout waiting for PCI detach, proceeding anyway")

# Get DPU configuration
dpu_ip = None
try:
dpu_ip = get_dpu_ip(self._config_db, dpu_name)
port = get_dpu_gnmi_port(self._config_db, dpu_name)
dpu_ip = get_dpu_ip(config_db, dpu_name)
port = get_dpu_gnmi_port(config_db, dpu_name)
if not dpu_ip:
logger.log_error(f"{dpu_name}: IP not found in DHCP_SERVER_IPV4_PORT table (key: bridge-midplane|{dpu_name.lower()}), cannot proceed")
self._clear_halt_flag(dpu_name)
Expand All @@ -206,16 +213,17 @@ def _handle_transition(self, dpu_name: str, transition_type: str) -> bool:

return reboot_successful

def _wait_for_gnoi_halt_in_progress(self, dpu_name: str) -> bool:
def _wait_for_gnoi_halt_in_progress(self, dpu_name: str, state_db=None) -> bool:
"""
Poll for gnoi_halt_in_progress flag in STATE_DB CHASSIS_MODULE_TABLE.
This flag is set by the platform after completing PCI detach.
"""
state_db = state_db if state_db is not None else self._db
deadline = time.monotonic() + _get_halt_timeout()

while time.monotonic() < deadline:
try:
table = swsscommon.Table(self._db, "CHASSIS_MODULE_TABLE")
table = swsscommon.Table(state_db, "CHASSIS_MODULE_TABLE")
(status, fvs) = table.get(dpu_name)

if status:
Expand Down Expand Up @@ -378,7 +386,15 @@ def main():
# Wrapper to clean up after transition
def handle_and_cleanup(dpu):
try:
reboot_handler._handle_transition(dpu, "shutdown")
# Per-thread DB connections: a redis connection is a
# single non-thread-safe socket, so sharing one lets
# concurrent DPU reads cross their IP/port values.
thread_config_db = daemon_base.db_connect("CONFIG_DB")
thread_state_db = daemon_base.db_connect("STATE_DB")
Comment thread
vvolam marked this conversation as resolved.

@vvolam vvolam Jul 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (optional, non-blocking): each DPU shutdown opens two fresh redis connections here that live only for this transition. swsscommon.DBConnector has no explicit close()/disconnect(), so they're reclaimed only when
handle_and_cleanup() returns and these locals go out of scope. That's fine given how infrequently DPUs shut down(≤8 per chassis), but good to fix this issue. Otherwise the PR looks good to me

@gpunathilell gpunathilell Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vvolam similar handling is done for all the other daemons in pmon (chassisd/xcvrd) etc. The disconnect is not called explicitly and handled by the garbage collection when it goes out of scope, the function only stays in scope until the shutdown happens, and then we can close it

reboot_handler._handle_transition(
dpu,
config_db=thread_config_db,
state_db=thread_state_db)
logger.log_info(f"{dpu}: Transition thread completed successfully")
except Exception as e:
logger.log_error(f"{dpu}: Transition thread failed with exception: {e}")
Expand Down
75 changes: 66 additions & 9 deletions tests/gnoi_shutdown_daemon_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def test_handle_transition_success(self, mock_monotonic, mock_sleep, mock_execut

with patch('gnoi_shutdown_daemon.swsscommon.Table', return_value=mock_table):
handler = gnoi_shutdown_daemon.GnoiRebootHandler(mock_db, mock_config_db, mock_chassis)
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

self.assertTrue(result)
mock_module.clear_module_gnoi_halt_in_progress.assert_called_once()
Expand Down Expand Up @@ -251,7 +251,7 @@ def test_handle_transition_gnoi_halt_timeout(self, mock_execute_command, mock_mo

with patch('gnoi_shutdown_daemon.swsscommon.Table', return_value=mock_table):
handler = gnoi_shutdown_daemon.GnoiRebootHandler(mock_db, mock_config_db, mock_chassis)
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should still succeed - code proceeds anyway after timeout warning
self.assertTrue(result)
Expand Down Expand Up @@ -301,7 +301,7 @@ def test_handle_transition_ip_failure(self, mock_get_gnmi_port, mock_get_dpu_ip,
# Mock _wait_for_gnoi_halt_in_progress to return immediately to prevent hanging
handler._wait_for_gnoi_halt_in_progress = MagicMock(return_value=True)

result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

self.assertFalse(result)
# Verify that clear_module_gnoi_halt_in_progress was called
Expand Down Expand Up @@ -523,7 +523,7 @@ def test_handle_transition_config_exception(self, mock_get_port, mock_get_ip, mo
# Mock _wait_for_gnoi_halt_in_progress to return immediately to prevent hanging
handler._wait_for_gnoi_halt_in_progress = MagicMock(return_value=True)

result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

self.assertFalse(result)
# Verify that clear_module_gnoi_halt_in_progress was called
Expand Down Expand Up @@ -603,7 +603,7 @@ def test_handle_transition_dpu_already_offline(self):
mock_chassis.get_module.return_value = mock_module

handler = gnoi_shutdown_daemon.GnoiRebootHandler(mock_db, mock_config_db, mock_chassis)
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should return True (success) without attempting gNOI reboot
self.assertTrue(result)
Expand All @@ -623,7 +623,7 @@ def test_handle_transition_dpu_powered_down(self):
mock_chassis.get_module.return_value = mock_module

handler = gnoi_shutdown_daemon.GnoiRebootHandler(mock_db, mock_config_db, mock_chassis)
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should return True (success) without attempting gNOI reboot
self.assertTrue(result)
Expand Down Expand Up @@ -652,7 +652,7 @@ def test_handle_transition_dpu_fault_proceeds(self):

with patch('gnoi_shutdown_daemon.get_dpu_ip', return_value="10.0.0.1"), \
patch('gnoi_shutdown_daemon.get_dpu_gnmi_port', return_value="8080"):
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should proceed with shutdown for Fault state
self.assertTrue(result)
Expand All @@ -675,7 +675,7 @@ def test_handle_transition_dpu_offline_clear_halt_failure(self):
# Make _clear_halt_flag fail
handler._clear_halt_flag = MagicMock(return_value=False)

result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should return False since _clear_halt_flag failed
self.assertFalse(result)
Expand All @@ -700,13 +700,70 @@ def test_handle_transition_oper_status_check_exception(self):

with patch('gnoi_shutdown_daemon.get_dpu_ip', return_value="10.0.0.1"), \
patch('gnoi_shutdown_daemon.get_dpu_gnmi_port', return_value="8080"):
result = handler._handle_transition("DPU0", "shutdown")
result = handler._handle_transition("DPU0")

# Should proceed with shutdown despite oper_status check failure
self.assertTrue(result)
handler._wait_for_gnoi_halt_in_progress.assert_called_once()
handler._send_reboot_command.assert_called_once()

@patch('gnoi_shutdown_daemon.daemon_base.db_connect')
@patch('gnoi_shutdown_daemon.GnoiRebootHandler')
@patch('gnoi_shutdown_daemon.swsscommon.ConfigDBConnector')
@patch('threading.Thread')
def test_handle_and_cleanup_per_thread_connections(self, mock_thread, mock_config_db_connector_class, mock_gnoi_reboot_handler, mock_db_connect):
"""Test that handle_and_cleanup opens per-thread DB connections and passes them to _handle_transition."""
mock_state_db = MagicMock()
mock_config_db = MagicMock()
mock_thread_config_db = MagicMock()
mock_thread_state_db = MagicMock()

# main() calls db_connect("STATE_DB") then ("CONFIG_DB");
# handle_and_cleanup then calls ("CONFIG_DB") and ("STATE_DB") for its own thread.
mock_db_connect.side_effect = [mock_state_db, mock_config_db,
mock_thread_config_db, mock_thread_state_db]
mock_config_db.hget.return_value = "down"

mock_pubsub = MagicMock()
mock_pubsub.get_message.side_effect = [mock_message, KeyboardInterrupt]
mock_redis_client = MagicMock()
mock_redis_client.pubsub.return_value = mock_pubsub
mock_config_db_connector = MagicMock()
mock_config_db_connector.db_name = "CONFIG_DB"
mock_config_db_connector.get_redis_client.return_value = mock_redis_client
mock_config_db_connector_class.return_value = mock_config_db_connector

mock_handler_instance = MagicMock()
mock_gnoi_reboot_handler.return_value = mock_handler_instance

mock_platform_submodule = MagicMock()
mock_sonic_platform = MagicMock()
mock_sonic_platform.platform = mock_platform_submodule

with patch.dict('sys.modules', {
'sonic_platform': mock_sonic_platform,
'sonic_platform.platform': mock_platform_submodule
}):
with self.assertRaises(KeyboardInterrupt):
gnoi_shutdown_daemon.main()

# Capture the thread target and run it synchronously to exercise handle_and_cleanup
thread_call_kwargs = mock_thread.call_args.kwargs
target_fn = thread_call_kwargs['target']
target_args = thread_call_kwargs['args']
target_fn(*target_args)

# Verify per-thread connections were opened
mock_db_connect.assert_any_call("CONFIG_DB")
mock_db_connect.assert_any_call("STATE_DB")

# Verify _handle_transition received the per-thread connections, not the shared ones
mock_handler_instance._handle_transition.assert_called_once_with(
"DPU0",
config_db=mock_thread_config_db,
state_db=mock_thread_state_db,
)


if __name__ == '__main__':
unittest.main()
Loading