Merge pull request #69 from stealthrocket/error-info

Preserve exception details

Author: chriso

src/dispatch/proto.py

@@ -2,6 +2,7 @@
 
 import pickle
 from dataclasses import dataclass
+from types import TracebackType
 from typing import Any
 
 import google.protobuf.any_pb2
@@ -280,13 +281,20 @@ class Error:
     Output.
     """
 
-    def __init__(self, status: Status, type: str | None, message: str | None):
+    def __init__(
+        self,
+        status: Status,
+        type: str | None,
+        message: str | None,
+        value: Exception | None = None,
+    ):
         """Create a new Error.
 
         Args:
             status: categorization of the error.
             type: arbitrary string, used for humans. Optional.
             message: arbitrary message. Optional.
+            value: arbitrary exception from which the error is derived. Optional.
 
         Raises:
             ValueError: Neither type or message was provided or status is
@@ -300,6 +308,7 @@ def __init__(self, status: Status, type: str | None, message: str | None):
         self.type = type
         self.message = message
         self.status = status
+        self.value = value
 
     @classmethod
     def from_exception(cls, ex: Exception, status: Status | None = None) -> Error:
@@ -313,18 +322,31 @@ def from_exception(cls, ex: Exception, status: Status | None = None) -> Error:
         if status is None:
             status = status_for_error(ex)
 
-        return Error(status, ex.__class__.__qualname__, str(ex))
+        return Error(status, ex.__class__.__qualname__, str(ex), ex)
 
     def to_exception(self) -> Exception:
-        # TODO: use correct error type
-        return RuntimeError(self.message)
+        """Returns an equivalent exception."""
+        if self.value is not None:
+            return self.value
+
+        g = globals()
+        try:
+            assert isinstance(self.type, str)
+            cls = g[self.type]
+            assert issubclass(cls, Exception)
+        except (KeyError, AssertionError):
+            return RuntimeError(self.message)
+        else:
+            return cls(self.message)
 
     @classmethod
     def _from_proto(cls, proto: error_pb.Error) -> Error:
-        return cls(Status.UNSPECIFIED, proto.type, proto.message)
+        value = pickle.loads(proto.value) if proto.value else None
+        return cls(Status.UNSPECIFIED, proto.type, proto.message, value)
 
     def _as_proto(self) -> error_pb.Error:
-        return error_pb.Error(type=self.type, message=self.message)
+        value = pickle.dumps(self.value) if self.value else None
+        return error_pb.Error(type=self.type, message=self.message, value=value)
 
 
 def _any_unpickle(any: google.protobuf.any_pb2.Any) -> Any:

src/dispatch/scheduler.py

@@ -222,7 +222,7 @@ def _run(self, input: Input) -> Output:
                     future = owner.result
                     assert future is not None
                 except (KeyError, AssertionError):
-                    logger.warning("skipping unexpected call result %s", cr)
+                    logger.warning("discarding unexpected call result %s", cr)
                     continue
 
                 logger.debug("dispatching %s to %s", call_result, owner)
@@ -254,7 +254,6 @@ def _run(self, input: Input) -> Output:
                     coroutine_id=coroutine.id, value=e.value
                 )
             except Exception as e:
-                raise
                 coroutine_result = CoroutineResult(coroutine_id=coroutine.id, error=e)
 
             # Handle coroutines that return or raise.
@@ -266,23 +265,27 @@ def _run(self, input: Input) -> Output:
 
                 # If this is the main coroutine, we're done.
                 if coroutine.parent_id is None:
-                    assert len(state.suspended) == 0
+                    for suspended in state.suspended.values():
+                        suspended.coroutine.close()
                     if coroutine_result.error is not None:
                         return Output.error(
                             Error.from_exception(coroutine_result.error)
                         )
                     return Output.value(coroutine_result.value)
 
                 # Otherwise, notify the parent of the result.
-                assert coroutine.parent_id in state.suspended
-                parent = state.suspended[coroutine.parent_id]
-                assert parent.result is not None
-                future = parent.result
-                future.add(coroutine_result)
-                if future.ready():
-                    state.ready.insert(0, parent)
-                    del state.suspended[parent.id]
-                    logger.debug("parent %s is now ready", parent)
+                try:
+                    parent = state.suspended[coroutine.parent_id]
+                    future = parent.result
+                    assert future is not None
+                except (KeyError, AssertionError):
+                    logger.warning("discarding %s", coroutine_result)
+                else:
+                    future.add(coroutine_result)
+                    if future.ready():
+                        state.ready.insert(0, parent)
+                        del state.suspended[parent.id]
+                        logger.debug("parent %s is now ready", parent)
                 continue
 
             # Handle coroutines that yield.
@@ -343,6 +346,11 @@ def _run(self, input: Input) -> Output:
             logger.exception("state could not be serialized")
             return Output.error(Error.from_exception(e, status=Status.PERMANENT_ERROR))
 
+        # Close coroutines before yielding.
+        for suspended in state.suspended.values():
+            suspended.coroutine.close()
+        state.suspended = {}
+
         # Yield to Dispatch.
         logger.debug(
             "yielding to Dispatch with %d call(s) and %d bytes of state",

tests/dispatch/test_scheduler.py

@@ -31,6 +31,11 @@ async def call_sequentially(*functions):
     return results
 
 
+@durable
+async def raises_error():
+    raise ValueError("oops")
+
+
 class TestOneShotScheduler(unittest.TestCase):
     def test_main_return(self):
         @durable
@@ -43,10 +48,10 @@ async def main():
     def test_main_raise(self):
         @durable
         async def main():
-            raise RuntimeError("oops")
+            raise ValueError("oops")
 
         output = self.start(main)
-        self.assert_exit_result_error(output, RuntimeError, "oops")
+        self.assert_exit_result_error(output, ValueError, "oops")
 
     def test_main_args(self):
         @durable
@@ -215,6 +220,14 @@ async def main():
 
         self.assertEqual(len(correlation_ids), 8)
 
+    def test_raise_indirect(self):
+        @durable
+        async def main():
+            return await gather(call_one("a"), raises_error())
+
+        output = self.start(main)
+        self.assert_exit_result_error(output, ValueError, "oops")
+
     def start(self, main: Callable, *args: Any, **kwargs: Any) -> Output:
         input = Input.from_input_arguments(main.__qualname__, *args, **kwargs)
         return OneShotScheduler(main).run(input)
@@ -258,6 +271,7 @@ def assert_exit_result_error(
         self.assertEqual(error.__class__, expect)
         if message is not None:
             self.assertEqual(str(error), message)
+        return error
 
     def assert_poll(self, output: Output) -> poll_pb.Poll:
         response = output._message