feat: utility function to persist critical error (#222)

Author: adarsh0728

pynumaflow/_constants.py

@@ -7,9 +7,11 @@
 SIDE_INPUT_DIR_PATH = "/var/numaflow/side-inputs"
 ENV_UD_CONTAINER_TYPE = "NUMAFLOW_UD_CONTAINER_TYPE"
 
-# Get container type from env var, default to unknown-container
+# Error Constants
+RUNTIME_APPLICATION_ERRORS_PATH = "/var/numaflow/runtime/application-errors"
+CURRENT_CRITICAL_ERROR_FILE = "current-udf.json"
+INTERNAL_ERROR_CODE = "Internal error"
 CONTAINER_TYPE = os.getenv(ENV_UD_CONTAINER_TYPE, "unknown-container")
-# UDF exception error string with container type
 ERR_UDF_EXCEPTION_STRING = f"UDF_EXECUTION_ERROR({CONTAINER_TYPE})"
 
 # Socket configs

pynumaflow/errors/__init__.py

@@ -0,0 +1,3 @@
+from pynumaflow.errors.errors import persist_critical_error
+
+__all__ = ["persist_critical_error"]

pynumaflow/errors/_dtypes.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class _RuntimeErrorEntry:
+    """Represents a runtime error entry to be persisted."""
+
+    container: str
+    timestamp: int
+    code: str
+    message: str
+    details: str
+
+    def to_dict(self) -> dict:
+        """Converts the dataclass instance to a dictionary."""
+        return asdict(self)

pynumaflow/errors/errors.py

@@ -0,0 +1,82 @@
+import os
+import json
+import threading
+import time
+from pynumaflow._constants import (
+    CONTAINER_TYPE,
+    RUNTIME_APPLICATION_ERRORS_PATH,
+    CURRENT_CRITICAL_ERROR_FILE,
+    INTERNAL_ERROR_CODE,
+)
+from pynumaflow.errors._dtypes import _RuntimeErrorEntry
+from typing import Union
+
+
+class _PersistErrorOnce:
+    """Ensures that the persist_critical_error function is executed only once."""
+
+    def __init__(self):
+        self.done = False
+        self.lock = threading.Lock()
+
+    def execute(self, func, *args, **kwargs):
+        with self.lock:
+            if self.done:
+                raise RuntimeError("Persist critical error function has already been executed.")
+            self.done = True
+            return func(*args, **kwargs)
+
+
+_persist_error_once = _PersistErrorOnce()
+
+
+def persist_critical_error(
+    error_code: str, error_message: str, error_details: str
+) -> Union[RuntimeError, None]:
+    """
+    Persists a critical error to a file. This function will only execute once.
+    Logs the error if persisting to the file fails.
+    Returns None if successful, or raises RuntimeError if already executed.
+    """
+    try:
+        _persist_error_once.execute(
+            _persist_critical_error_to_file,
+            error_code,
+            error_message,
+            error_details,
+            RUNTIME_APPLICATION_ERRORS_PATH,
+        )
+    except RuntimeError as e:
+        return e
+    except Exception as e:
+        print(f"Error in persisting critical error: {e}")
+    return None
+
+
+def _persist_critical_error_to_file(
+    error_code: str, error_message: str, error_details: str, dir_path: str
+):
+    """Internal function to persist a critical error to a file."""
+
+    os.makedirs(dir_path, mode=0o777, exist_ok=True)
+    container_dir = os.path.join(dir_path, CONTAINER_TYPE)
+    os.makedirs(container_dir, mode=0o777, exist_ok=True)
+
+    current_file_path = os.path.join(container_dir, CURRENT_CRITICAL_ERROR_FILE)
+    error_code = error_code or INTERNAL_ERROR_CODE
+    current_timestamp = int(time.time())
+
+    runtime_error_entry = _RuntimeErrorEntry(
+        container=CONTAINER_TYPE,
+        timestamp=current_timestamp,
+        code=error_code,
+        message=error_message,
+        details=error_details,
+    )
+
+    with open(current_file_path, "w") as f:
+        json.dump(runtime_error_entry.to_dict(), f)
+
+    final_file_name = f"{current_timestamp}-udf.json"
+    final_file_path = os.path.join(container_dir, final_file_name)
+    os.rename(current_file_path, final_file_path)

tests/errors/__init__.py

@@ -0,0 +1,72 @@
+import unittest
+from pynumaflow.errors._dtypes import _RuntimeErrorEntry
+
+
+class TestRuntimeErrorEntry(unittest.TestCase):
+    def test_runtime_error_entry_initialization(self):
+        """
+        Test that _RuntimeErrorEntry initializes correctly with given values.
+        """
+        container = "test-container"
+        timestamp = 1680700000
+        code = "500"
+        message = "Test error message"
+        details = "Test error details"
+
+        error_entry = _RuntimeErrorEntry(container, timestamp, code, message, details)
+
+        self.assertEqual(error_entry.container, container)
+        self.assertEqual(error_entry.timestamp, timestamp)
+        self.assertEqual(error_entry.code, code)
+        self.assertEqual(error_entry.message, message)
+        self.assertEqual(error_entry.details, details)
+
+    def test_runtime_error_entry_to_dict(self):
+        """
+        Test that _RuntimeErrorEntry converts to a dictionary correctly.
+        """
+        container = "test-container"
+        timestamp = 1680700000
+        code = "500"
+        message = "Test error message"
+        details = "Test error details"
+
+        error_entry = _RuntimeErrorEntry(container, timestamp, code, message, details)
+        error_dict = error_entry.to_dict()
+
+        expected_dict = {
+            "container": container,
+            "timestamp": timestamp,
+            "code": code,
+            "message": message,
+            "details": details,
+        }
+
+        self.assertEqual(error_dict, expected_dict)
+
+    def test_runtime_error_entry_empty_values(self):
+        """
+        Test that _RuntimeErrorEntry handles empty values correctly.
+        """
+        container = ""
+        timestamp = 0
+        code = ""
+        message = ""
+        details = ""
+
+        error_entry = _RuntimeErrorEntry(container, timestamp, code, message, details)
+        error_dict = error_entry.to_dict()
+
+        expected_dict = {
+            "container": container,
+            "timestamp": timestamp,
+            "code": code,
+            "message": message,
+            "details": details,
+        }
+
+        self.assertEqual(error_dict, expected_dict)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/errors/test_dtypes.py

@@ -0,0 +1,134 @@
+import os
+import json
+import shutil
+import threading
+import unittest
+from pynumaflow.errors.errors import persist_critical_error, _persist_error_once
+from pynumaflow.errors.errors import _persist_critical_error_to_file
+from pynumaflow._constants import CONTAINER_TYPE, INTERNAL_ERROR_CODE
+
+
+class TestErrorPersistence(unittest.TestCase):
+    def setUp(self):
+        """
+        Set up temporary directories for tests.
+        """
+        self.test_dirs = ["/tmp/test_error_dir", "/tmp/test_dir"]
+
+    def tearDown(self):
+        """
+        Clean up temporary directories after tests.
+        """
+        for dir_path in self.test_dirs:
+            if os.path.exists(dir_path):
+                shutil.rmtree(dir_path)
+
+    # Writes error details to a JSON file
+    def test_writes_error_details_to_json_file(self):
+        """
+        Test that _persist_critical_error_to_file writes error details to a JSON file.
+        """
+
+        dir_path = self.test_dirs[0]
+
+        error_code = "500"
+        error_message = "Server Error"
+        error_details = "An unexpected error occurred."
+
+        _persist_critical_error_to_file(error_code, error_message, error_details, dir_path)
+
+        container_dir = os.path.join(dir_path, CONTAINER_TYPE)
+        self.assertTrue(os.path.exists(container_dir))
+
+        # Debug: Check directory after the function call
+        print(f"After: {os.listdir(container_dir)}")
+
+        files = os.listdir(container_dir)
+        self.assertEqual(len(files), 1)
+
+        final_file_name = files[0]
+        final_file_path = os.path.join(container_dir, final_file_name)
+
+        with open(final_file_path) as f:
+            data = json.load(f)
+
+        self.assertEqual(data["code"], error_code)
+        self.assertEqual(data["message"], error_message)
+        self.assertEqual(data["details"], error_details)
+        self.assertEqual(data["container"], CONTAINER_TYPE)
+        self.assertTrue(isinstance(data["timestamp"], int))
+
+    # Uses default error code if none provided
+    def test_uses_default_error_code_if_none_provided(self):
+        """
+        Test that _persist_critical_error_to_file uses the default error code if none is provided.
+        """
+        dir_path = self.test_dirs[1]
+
+        _persist_critical_error_to_file("", "Error Message", "Error Details", dir_path)
+
+        container_dir = os.path.join(dir_path, "unknown-container")
+        self.assertTrue(os.path.exists(container_dir))
+
+        files = os.listdir(container_dir)
+        self.assertEqual(len(files), 1)
+
+        with open(os.path.join(container_dir, files[0])) as f:
+            error_data = json.load(f)
+            self.assertEqual(error_data["code"], INTERNAL_ERROR_CODE)
+
+    def test_persist_critical_error_all_threads_fail(self):
+        """
+        Test that all threads fail when persist_critical_error is executed after the first call.
+        """
+        error_code = "testCode"
+        error_message = "testMessage"
+        error_details = "testDetails"
+
+        # Set `done` to True to simulate that the critical error has already been persisted
+        _persist_error_once.done = True
+
+        try:
+            # Set up threading
+            num_threads = 10
+            errors = []
+            lock = threading.Lock()
+
+            def thread_func():
+                nonlocal errors
+                result = persist_critical_error(error_code, error_message, error_details)
+                with lock:
+                    errors.append(result)
+
+            # Create and start threads
+            threads = []
+            for _ in range(num_threads):
+                thread = threading.Thread(target=thread_func)
+                threads.append(thread)
+                thread.start()
+
+            # Wait for all threads to complete
+            for thread in threads:
+                thread.join()
+
+            # Count the number of failures
+            fail_count = sum(
+                1
+                for error in errors
+                if error is not None
+                and "Persist critical error function has already been executed" in str(error)
+            )
+
+            # Assert that all threads failed
+            self.assertEqual(
+                fail_count,
+                num_threads,
+                f"Expected all {num_threads} threads to fail, but only {fail_count} failed",
+            )
+        finally:
+            # Revert `done` back to False after the test
+            _persist_error_once.done = False
+
+
+if __name__ == "__main__":
+    unittest.main()