Fix stacktrace propagation in python (#171)

Author: slinkydeveloper

python/restate/server_context.py

@@ -427,7 +427,7 @@ async def enter(self):
                 cause = cause.__cause__
             else:
                 # nothing interesting found, treat as unexpected exception
-                stacktrace = "\n".join(traceback.format_exception(e))
+                stacktrace = "".join(traceback.format_exception(e))
                 self.vm.notify_error(repr(e), stacktrace)
                 raise
         finally:
@@ -680,7 +680,7 @@ async def create_run_coroutine(
         except Exception as e:
             end = time.time()
             attempt_duration = int((end - start) * 1000)
-            failure = Failure(code=500, message=str(e))
+            failure = Failure(code=500, message=repr(e), stacktrace="".join(traceback.format_exception(e)))
             max_duration_ms = None if max_duration is None else int(max_duration.total_seconds() * 1000)
             initial_retry_interval_ms = (
                 None if initial_retry_interval is None else int(initial_retry_interval.total_seconds() * 1000)

python/restate/vm.py

@@ -71,6 +71,7 @@ class Failure:
 
     code: int
     message: str
+    stacktrace: typing.Optional[str] = None
 
 
 @dataclass
@@ -433,7 +434,7 @@ def propose_run_completion_transient(
         Exit a side effect with a transient Error.
         This requires a retry policy to be provided.
         """
-        py_failure = PyFailure(failure.code, failure.message)
+        py_failure = PyFailure(failure.code, failure.message, failure.stacktrace)
         py_config = PyExponentialRetryConfig(
             config.initial_interval,
             config.max_attempts,

src/lib.rs

@@ -96,13 +96,20 @@ struct PyFailure {
     code: u16,
     #[pyo3(get, set)]
     message: String,
+    #[pyo3(get, set)]
+    stacktrace: Option<String>,
 }
 
 #[pymethods]
 impl PyFailure {
     #[new]
-    fn new(code: u16, message: String) -> PyFailure {
-        Self { code, message }
+    #[pyo3(signature = (code, message, stacktrace=None))]
+    fn new(code: u16, message: String, stacktrace: Option<String>) -> PyFailure {
+        Self {
+            code,
+            message,
+            stacktrace,
+        }
     }
 }
 
@@ -173,6 +180,7 @@ impl From<TerminalFailure> for PyFailure {
         PyFailure {
             code: value.code,
             message: value.message,
+            stacktrace: None,
         }
     }
 }
@@ -188,8 +196,18 @@ impl From<PyFailure> for TerminalFailure {
 }
 
 impl From<PyFailure> for Error {
-    fn from(value: PyFailure) -> Self {
-        Self::new(value.code, value.message)
+    fn from(
+        PyFailure {
+            code,
+            message,
+            stacktrace,
+        }: PyFailure,
+    ) -> Self {
+        let mut e = Self::new(code, message);
+        if let Some(stacktrace) = stacktrace {
+            e = e.with_stacktrace(stacktrace);
+        }
+        e
     }
 }
 
@@ -501,7 +519,8 @@ impl PyVM {
                         .map(Into::into)
                         .collect(),
                 },
-                buffer.as_bytes().to_vec().into(),Default::default()
+                buffer.as_bytes().to_vec().into(),
+                Default::default(),
             )
             .map(Into::into)
             .map_err(Into::into)
@@ -539,8 +558,8 @@ impl PyVM {
                         .duration_since(SystemTime::UNIX_EPOCH)
                         .expect("Duration since unix epoch cannot fail")
                         + Duration::from_millis(millis)
-                })
-                , Default::default()
+                }),
+                Default::default(),
             )
             .map(|s| s.invocation_id_notification_handle.into())
             .map_err(Into::into)
@@ -566,7 +585,7 @@ impl PyVM {
             .sys_complete_awakeable(
                 id,
                 NonEmptyValue::Success(buffer.as_bytes().to_vec().into()),
-                Default::default()
+                Default::default(),
             )
             .map_err(Into::into)
     }
@@ -613,7 +632,8 @@ impl PyVM {
             .vm
             .sys_complete_promise(
                 key,
-                NonEmptyValue::Success(buffer.as_bytes().to_vec().into()),Default::default()
+                NonEmptyValue::Success(buffer.as_bytes().to_vec().into()),
+                Default::default(),
             )
             .map(Into::into)
             .map_err(Into::into)
@@ -626,7 +646,11 @@ impl PyVM {
     ) -> Result<PyNotificationHandle, PyVMError> {
         self_
             .vm
-            .sys_complete_promise(key, NonEmptyValue::Failure(value.into()),Default::default())
+            .sys_complete_promise(
+                key,
+                NonEmptyValue::Failure(value.into()),
+                Default::default(),
+            )
             .map(Into::into)
             .map_err(Into::into)
     }
@@ -701,7 +725,10 @@ impl PyVM {
     ) -> Result<(), PyVMError> {
         self_
             .vm
-            .sys_write_output(NonEmptyValue::Success(buffer.as_bytes().to_vec().into()),Default::default())
+            .sys_write_output(
+                NonEmptyValue::Success(buffer.as_bytes().to_vec().into()),
+                Default::default(),
+            )
             .map_err(Into::into)
     }
 
@@ -711,7 +738,7 @@ impl PyVM {
     ) -> Result<(), PyVMError> {
         self_
             .vm
-            .sys_write_output(NonEmptyValue::Failure(value.into()),Default::default())
+            .sys_write_output(NonEmptyValue::Failure(value.into()), Default::default())
             .map_err(Into::into)
     }