Fix co_await argument dependent lookup for TFuture<T>

commit_hash:d3da906e0b3f303ead8bf1257a3cdceed2778d90

4 files changed, 84 insertions, 16 deletions

diff --git a/library/cpp/threading/future/core/coroutine_traits.h b/library/cpp/threading/future/core/coroutine_traits.h
index de08c377b16..14c5539eae0 100644
--- a/library/cpp/threading/future/core/coroutine_traits.h
+++ b/library/cpp/threading/future/core/coroutine_traits.h
@@ -126,23 +126,19 @@ namespace NThreading {
     template <typename T>
     using TExtractingFutureAwaitable = TFutureAwaitable<T, true>;
 
-} // namespace NThreading
-
-template <typename T>
-auto operator co_await(const NThreading::TFuture<T>& future) noexcept {
-    return NThreading::TFutureAwaitable{future};
-}
-
-template <typename T>
-auto operator co_await(NThreading::TFuture<T>&& future) noexcept {
-    // Not TExtractongFutureAwaitable, because TFuture works like std::shared_future.
-    // auto value = co_await GetCachedFuture();
-    // If GetCachedFuture stores a future in some cache and returns its copies,
-    // then subsequent uses of co_await will return a moved-from value.
-    return NThreading::TFutureAwaitable{std::move(future)};
-}
+    template <typename T>
+    auto operator co_await(const NThreading::TFuture<T>& future) noexcept {
+        return NThreading::TFutureAwaitable{future};
+    }
 
-namespace NThreading {
+    template <typename T>
+    auto operator co_await(NThreading::TFuture<T>&& future) noexcept {
+        // Not TExtractongFutureAwaitable, because TFuture works like std::shared_future.
+        // auto value = co_await GetCachedFuture();
+        // If GetCachedFuture stores a future in some cache and returns its copies,
+        // then subsequent uses of co_await will return a moved-from value.
+        return NThreading::TFutureAwaitable{std::move(future)};
+    }
 
     template <typename T>
     auto AsAwaitable(const NThreading::TFuture<T>& fut) noexcept {
diff --git a/library/cpp/threading/future/ut_gtest/coroutine_traits_adl_ut.cpp b/library/cpp/threading/future/ut_gtest/coroutine_traits_adl_ut.cpp
new file mode 100644
index 00000000000..317dcb73c03
--- /dev/null
+++ b/library/cpp/threading/future/ut_gtest/coroutine_traits_adl_ut.cpp
@@ -0,0 +1,24 @@
+#include "simple_task.h"
+
+#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>
+
+TEST(TestFutureTraits, ArgumentDependentLookup) {
+    TVector<TString> result;
+    NThreading::TPromise<void> promise = NThreading::NewPromise<void>();
+    NThreading::TFuture<void> future = promise.GetFuture();
+
+    NSimpleTestTask::RunAwaitable(result, std::move(future));
+
+    promise.SetValue();
+
+    EXPECT_THAT(
+        result,
+        ::testing::ContainerEq(
+            TVector<TString>({
+                "before co_await",
+                "after co_await",
+            })));
+}
diff --git a/library/cpp/threading/future/ut_gtest/simple_task.h b/library/cpp/threading/future/ut_gtest/simple_task.h
new file mode 100644
index 00000000000..8291de2dba0
--- /dev/null
+++ b/library/cpp/threading/future/ut_gtest/simple_task.h
@@ -0,0 +1,47 @@
+#pragma once
+
+#include <util/generic/string.h>
+#include <util/generic/vector.h>
+
+#include <coroutine>
+#include <utility>
+
+// Note: this namespace must be separate from NThreading and global namespace
+namespace NSimpleTestTask {
+
+    /**
+     * The simplest possible coroutine type for testing
+     */
+    struct TSimpleTask {
+        struct promise_type {
+            TSimpleTask get_return_object() noexcept {
+                return {};
+            }
+            std::suspend_never initial_suspend() noexcept {
+                return {};
+            }
+            std::suspend_never final_suspend() noexcept {
+                return {};
+            }
+            void unhandled_exception() {
+                throw;
+            }
+            void return_void() noexcept {
+            }
+        };
+    };
+
+    /**
+     * This coroutine awaits a generic awaitable
+     *
+     * When operator co_await for awaitables is not defined before this header
+     * is included it will be searched using argument dependent lookup.
+     */
+    template <class TAwaitable>
+    TSimpleTask RunAwaitable(TVector<TString>& result, TAwaitable&& awaitable) {
+        result.push_back("before co_await");
+        co_await std::forward<TAwaitable>(awaitable);
+        result.push_back("after co_await");
+    }
+
+} // namespace NSimpleTestTask
diff --git a/library/cpp/threading/future/ut_gtest/ya.make b/library/cpp/threading/future/ut_gtest/ya.make
index 84f7ebe2e0e..509180f9918 100644
--- a/library/cpp/threading/future/ut_gtest/ya.make
+++ b/library/cpp/threading/future/ut_gtest/ya.make
@@ -5,6 +5,7 @@ PEERDIR(
 )
 
 SRCS(
+    coroutine_traits_adl_ut.cpp
     coroutine_traits_ut.cpp
 )