fix(rollback): resolve signal handler race condition by clearing context after execution (#1245)

- Fixed race condition in signal handler
- Cleared context after execution instead of before
- Improved rollback stability

Signed-off-by: Nitesh <nitesh@example.com>
Co-authored-by: Nitesh <nitesh@example.com>
Co-authored-by: Paige Patton <64206430+paigerube14@users.noreply.github.com>

krkn/prometheus/collector.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+#
 # Copyright 2025 The Krkn Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 from __future__ import annotations
 
 import datetime

krkn/resiliency/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """krkn.resiliency package public interface."""
 
 from .resiliency import Resiliency  # noqa: F401

krkn/resiliency/resiliency.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+#
 # Copyright 2025 The Krkn Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Resiliency evaluation orchestrator for Krkn chaos runs.
 
 This module provides the `Resiliency` class which loads the canonical

krkn/resiliency/score.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+#
 # Copyright 2025 The Krkn Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 from __future__ import annotations
 
 from typing import Dict, List, Tuple

krkn/rollback/signal.py

@@ -58,13 +58,19 @@ class SignalHandler:
             logger.warning(f"Signal {signal_name} received without complete context, skipping rollback.")
             return
 
-        # Clear the context for the next signal, as another signal may arrive before the rollback completes.
-        # This ensures that the rollback is performed only once.
-        cls._set_context(None, None, telemetry_ocp)
-        
-        # Perform rollback
-        logger.info(f"Performing rollback for signal {signal_name} with run_uuid={run_uuid}, scenario_type={scenario_type}")
-        execute_rollback_version_files(telemetry_ocp, run_uuid, scenario_type)
+        # Prevent concurrent rollback execution with a lock
+        if not cls._signal_lock.acquire(blocking=False):
+            logger.warning(f"Signal {signal_name} received but rollback already in progress, skipping.")
+            return
+
+        try:
+            # Perform rollback
+            logger.info(f"Performing rollback for signal {signal_name} with run_uuid={run_uuid}, scenario_type={scenario_type}")
+            execute_rollback_version_files(telemetry_ocp, run_uuid, scenario_type)
+        finally:
+            # Always clear context and release lock after rollback completes or fails
+            cls._set_context(None, None, telemetry_ocp)
+            cls._signal_lock.release()
         
         # Call original handler if it exists
         if signum not in cls._original_handlers:

tests/test_prometheus_collector.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 """
 Tests for krkn.prometheus.collector module.
 

tests/test_resiliency.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 """
 Tests for krkn.resiliency.resiliency module.
 

tests/test_resiliency_score.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 """
 Tests for krkn.resiliency.score module.