Allow to pass exception to the coroutine to be cancelled

3 files changed, 50 insertions, 0 deletions

diff --git a/library/cpp/coroutine/engine/coroutine_ut.cpp b/library/cpp/coroutine/engine/coroutine_ut.cpp
index 35b7d6353a..c8a1852f7c 100644
--- a/library/cpp/coroutine/engine/coroutine_ut.cpp
+++ b/library/cpp/coroutine/engine/coroutine_ut.cpp
@@ -47,6 +47,7 @@ class TCoroTest: public TTestBase {
     UNIT_TEST(TestUserEvent);
     UNIT_TEST(TestPause);
     UNIT_TEST(TestOverrideTime);
+    UNIT_TEST(TestCancelWithException);
     UNIT_TEST_SUITE_END();
 
 public:
@@ -79,6 +80,7 @@ public:
     void TestUserEvent();
     void TestPause();
     void TestOverrideTime();
+    void TestCancelWithException();
 };
 
 void TCoroTest::TestException() {
@@ -1005,4 +1007,38 @@ void TCoroTest::TestOverrideTime() {
 
     executor.Execute();
 }
+
+void TCoroTest::TestCancelWithException() {
+    TContExecutor exec(32000);
+
+    TString excText = "test exception";
+    THolder<std::exception> excep = MakeHolder<yexception>(yexception() << excText);
+    std::exception* excPtr = excep.Get();
+
+    exec.CreateOwned([&](TCont* cont){
+        TCont *cont1 = cont->Executor()->CreateOwned([&](TCont* c) {
+            int result = c->SleepD(TDuration::MilliSeconds(200).ToDeadLine());
+            UNIT_ASSERT_EQUAL(result, ECANCELED);
+            UNIT_ASSERT_EQUAL(c->Cancelled(), true);
+            THolder<std::exception> exc = c->TakeException();
+            UNIT_ASSERT_EQUAL(exc.Get(), excPtr);
+            UNIT_ASSERT_EQUAL(exc->what(), excText);
+            UNIT_ASSERT(dynamic_cast<yexception*>(exc.Get()) != nullptr);
+        }, "cancelExc");
+        cont1->Cancel(std::move(excep));
+
+        TCont* cont2 = cont->Executor()->CreateOwned([&](TCont* c) {
+            int result = c->SleepD(TDuration::MilliSeconds(200).ToDeadLine());
+            UNIT_ASSERT_EQUAL(result, ECANCELED);
+            UNIT_ASSERT_EQUAL(c->Cancelled(), true);
+            THolder<std::exception> exc = c->TakeException();
+            UNIT_ASSERT_EQUAL(exc.Get(), nullptr);
+        }, "cancelTwice");
+        cont2->Cancel();
+        THolder<std::exception> e = MakeHolder<yexception>(yexception() << "another exception");
+        cont2->Cancel(std::move(e));
+    }, "coro");
+
+    exec.Execute();
+}
 UNIT_TEST_SUITE_REGISTRATION(TCoroTest);
diff --git a/library/cpp/coroutine/engine/impl.cpp b/library/cpp/coroutine/engine/impl.cpp
index 7ae6f74051..08f358127d 100644
--- a/library/cpp/coroutine/engine/impl.cpp
+++ b/library/cpp/coroutine/engine/impl.cpp
@@ -106,6 +106,13 @@ void TCont::Cancel() noexcept {
     }
 }
 
+void TCont::Cancel(THolder<std::exception> exception) noexcept {
+    if (!Cancelled()) {
+        Exception_ = std::move(exception);
+        Cancel();
+    }
+}
+
 void TCont::ReSchedule() noexcept {
     if (Cancelled()) {
         // Legacy code may expect a Cancelled coroutine to be scheduled without delay.
diff --git a/library/cpp/coroutine/engine/impl.h b/library/cpp/coroutine/engine/impl.h
index 3795fdcec1..ca523cdaf3 100644
--- a/library/cpp/coroutine/engine/impl.h
+++ b/library/cpp/coroutine/engine/impl.h
@@ -86,6 +86,7 @@ public:
     bool IAmRunning() const noexcept;
 
     void Cancel() noexcept;
+    void Cancel(THolder<std::exception> exception) noexcept;
 
     bool Cancelled() const noexcept {
         return Cancelled_;
@@ -108,6 +109,10 @@ public:
         Trampoline_.SwitchTo(ctx);
     }
 
+    THolder<std::exception> TakeException() noexcept {
+        return std::move(Exception_);
+    }
+
 private:
     void Terminate();
 
@@ -122,6 +127,8 @@ private:
     TIntrusiveList<TJoinWait> Waiters_;
     bool Cancelled_ = false;
     bool Scheduled_ = false;
+
+    THolder<std::exception> Exception_;
 };
 
 TCont* RunningCont();