From e43e1b272d0b7bc3a0eeb2ec4d487188f839ff10 Mon Sep 17 00:00:00 2001
From: robot-piglet <robot-piglet@yandex-team.com>
Date: Thu, 1 Aug 2024 16:52:56 +0300
Subject: Intermediate changes

---
 .../future/ut_gtest/coroutine_traits_ut.cpp        | 189 +++++++++++++++++++++
 library/cpp/threading/future/ut_gtest/ya.make      |  11 ++
 library/cpp/threading/future/ya.make               |   1 +
 3 files changed, 201 insertions(+)
 create mode 100644 library/cpp/threading/future/ut_gtest/coroutine_traits_ut.cpp
 create mode 100644 library/cpp/threading/future/ut_gtest/ya.make

(limited to 'library/cpp/threading')

diff --git a/library/cpp/threading/future/ut_gtest/coroutine_traits_ut.cpp b/library/cpp/threading/future/ut_gtest/coroutine_traits_ut.cpp
new file mode 100644
index 00000000000..a98e455c583
--- /dev/null
+++ b/library/cpp/threading/future/ut_gtest/coroutine_traits_ut.cpp
@@ -0,0 +1,189 @@
+#include <library/cpp/testing/gtest/gtest.h>
+#include <library/cpp/testing/gtest_extensions/gtest_extensions.h>
+#include <library/cpp/threading/future/future.h>
+#include <library/cpp/threading/future/core/coroutine_traits.h>
+
+#include <util/generic/vector.h>
+#include <util/generic/scope.h>
+#include <util/system/platform.h>
+
+
+TEST(TestFutureTraits, Simple) {
+    TVector<TString> result;
+
+    auto coroutine1 = [&result]() -> NThreading::TFuture<size_t> {
+        result.push_back("coroutine1");
+        co_return 1;
+    };
+
+    NThreading::TPromise<size_t> coroutine2SuspendPromise = NThreading::NewPromise<size_t>();
+    auto coroutine2 = [&result, coroutine2SuspendFuture = coroutine2SuspendPromise.GetFuture()]() -> NThreading::TFuture<size_t> {
+        result.push_back("coroutine2");
+
+        result.push_back("pre_coroutine2_suspend_future");
+        size_t futureResult = co_await coroutine2SuspendFuture;
+        result.push_back("post_coroutine2_suspend_future");
+
+        co_return 2 + futureResult;
+    };
+
+    auto coroutineAll = [&]() -> NThreading::TFuture<size_t> {
+        Y_DEFER {
+            result.push_back("coroutine_all_destroy");
+        };
+
+        result.push_back("pre_coroutine1");
+        size_t coroutine1Res = co_await coroutine1();
+        result.push_back("post_coroutine1");
+
+        result.push_back("pre_coroutine2");
+        size_t coroutine2Res = co_await coroutine2();
+        result.push_back("post_coroutine2");
+
+        co_return coroutine1Res + coroutine2Res;
+    };
+
+    NThreading::TFuture<size_t> coroutineAllFuture = coroutineAll();
+    EXPECT_FALSE(coroutineAllFuture.HasValue());
+    EXPECT_FALSE(coroutineAllFuture.HasException());
+    EXPECT_THAT(
+        result,
+        ::testing::ContainerEq(
+            TVector<TString>({
+                "pre_coroutine1",
+                "coroutine1",
+                "post_coroutine1",
+
+                "pre_coroutine2",
+                "coroutine2",
+                "pre_coroutine2_suspend_future"
+            })
+        )
+    );
+
+    coroutine2SuspendPromise.SetValue(3u);
+    EXPECT_TRUE(coroutineAllFuture.HasValue());
+    EXPECT_EQ(coroutineAllFuture.GetValue(), 6u);
+    EXPECT_THAT(
+        result,
+        ::testing::ContainerEq(
+            TVector<TString>({
+                "pre_coroutine1",
+                "coroutine1",
+                "post_coroutine1",
+
+                "pre_coroutine2",
+                "coroutine2",
+                "pre_coroutine2_suspend_future",
+                "post_coroutine2_suspend_future",
+                "post_coroutine2",
+
+                "coroutine_all_destroy"
+            })
+        )
+    );
+}
+
+TEST(TestFutureTraits, Exception) {
+    TVector<TString> result;
+
+    auto coroutine1 = [&result]() -> NThreading::TFuture<size_t> {
+        result.push_back("coroutine1");
+        co_return 1;
+    };
+
+    auto coroutine2 = [&result]() -> NThreading::TFuture<size_t> {
+        result.push_back("coroutine2");
+        ythrow yexception() << "coroutine2 exception";
+    };
+
+    auto coroutineAll = [&]() -> NThreading::TFuture<size_t> {
+        Y_DEFER {
+            result.push_back("coroutine_all_destroy");
+        };
+
+        result.push_back("pre_coroutine1");
+        size_t coroutine1Res = co_await coroutine1();
+        result.push_back("post_coroutine1");
+
+        result.push_back("pre_coroutine2");
+        size_t coroutine2Res = co_await coroutine2();
+        result.push_back("post_coroutine2");
+
+        co_return coroutine1Res + coroutine2Res;
+    };
+
+    EXPECT_THROW_MESSAGE_HAS_SUBSTR(
+        coroutineAll().GetValueSync(),
+        yexception,
+        "coroutine2 exception"
+    );
+    EXPECT_THAT(
+        result,
+        ::testing::ContainerEq(
+            TVector<TString>({
+                "pre_coroutine1",
+                "coroutine1",
+                "post_coroutine1",
+
+                "pre_coroutine2",
+                "coroutine2",
+
+                "coroutine_all_destroy"
+            })
+        )
+    );
+}
+
+TEST(TestFutureTraits, CrashOnExceptionInCoroutineHandlerResume) {
+    EXPECT_DEATH(
+        {
+            struct TBadPromise;
+            struct TBadCoroutine : std::coroutine_handle<TBadPromise> {
+                using promise_type = TBadPromise;
+            };
+
+            struct TBadPromise {
+                TBadCoroutine get_return_object() {
+                    return {TBadCoroutine::from_promise(*this)};
+                }
+
+                std::suspend_never initial_suspend() noexcept {
+                    return {};
+                }
+                std::suspend_never final_suspend() noexcept {
+                    return {};
+                }
+                void return_void() {
+                }
+                void unhandled_exception() {
+                    throw;
+                }
+            };
+
+            auto badCoroutine = []() -> TBadCoroutine {
+                ythrow yexception() << "bad coroutine exception";
+            };
+            // Sanity check
+            EXPECT_THROW_MESSAGE_HAS_SUBSTR(
+                badCoroutine(),
+                yexception,
+                "bad coroutine exception"
+            );
+
+            NThreading::TPromise<void> promise = NThreading::NewPromise<void>();
+            auto badCoroutineWithFutureAwait = [future = promise.GetFuture()]() -> TBadCoroutine {
+                co_await future;
+                ythrow yexception() << "bad coroutine with future await exception";
+            };
+
+            badCoroutineWithFutureAwait();
+            promise.SetValue();
+        },
+#if defined(_win_)
+        ".*"
+#else
+        "bad coroutine with future await exception"
+#endif
+    );
+}
diff --git a/library/cpp/threading/future/ut_gtest/ya.make b/library/cpp/threading/future/ut_gtest/ya.make
new file mode 100644
index 00000000000..84f7ebe2e0e
--- /dev/null
+++ b/library/cpp/threading/future/ut_gtest/ya.make
@@ -0,0 +1,11 @@
+GTEST()
+
+PEERDIR(
+    library/cpp/threading/future
+)
+
+SRCS(
+    coroutine_traits_ut.cpp
+)
+
+END()
diff --git a/library/cpp/threading/future/ya.make b/library/cpp/threading/future/ya.make
index 1d98569906c..77ee8b3c539 100644
--- a/library/cpp/threading/future/ya.make
+++ b/library/cpp/threading/future/ya.make
@@ -19,4 +19,5 @@ RECURSE(
     mt_ut
     perf
     ut
+    ut_gtest
 )
-- 
cgit v1.3