text: /*
   * Copyright (c) Facebook, Inc. and its affiliates.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #include <folly/Portability.h>
  
  #if FOLLY_HAS_COROUTINES
  
  #include <folly/executors/InlineExecutor.h>
  #include <folly/executors/ManualExecutor.h>
  #include <folly/experimental/coro/Baton.h>
  #include <folly/experimental/coro/BlockingWait.h>
  #include <folly/experimental/coro/Invoke.h>
  #include <folly/experimental/coro/Mutex.h>
  #include <folly/experimental/coro/SharedMutex.h>
  #include <folly/experimental/coro/Task.h>
  #include <folly/experimental/coro/detail/InlineTask.h>
  #include <folly/futures/Future.h>
  #include <folly/portability/GTest.h>
  
  #include <type_traits>
  
  using namespace folly;
  
  static_assert(
      std::is_same<
          folly::coro::semi_await_result_t<folly::coro::Task<void>>,
          void>::value,
      "");
  static_assert(
      std::is_same<
          folly::coro::semi_await_result_t<folly::coro::Task<int>>,
          int>::value,
      "");
  
  static_assert(
      std::is_same<
          folly::coro::semi_await_result_t<folly::coro::detail::InlineTask<void>>,
          void>::value,
      "");
  static_assert(
      std::is_same<
          folly::coro::semi_await_result_t<folly::coro::detail::InlineTask<int>>,
          int>::value,
      "");
  
  static_assert(
      std::is_same<folly::coro::semi_await_result_t<folly::coro::Baton&>, void>::
          value,
      "");
  static_assert(
      std::is_same<
          folly::coro::semi_await_result_t<decltype(
              std::declval<folly::coro::SharedMutex&>().co_scoped_lock_shared())>,
          folly::coro::SharedLock<folly::coro::SharedMutex>>::value,
      "");
  
  namespace {
  
  const RequestToken testToken1("corotest1");
  const RequestToken testToken2("corotest2");
  
  class TestRequestData : public RequestData {
   public:
    explicit TestRequestData(std::string key) noexcept : key_(std::move(key)) {}
  
    bool hasCallback() override {
      return false;
    }
  
    const std::string& key() const noexcept {
      return key_;
    }
  
   private:
    std::string key_;
  };
  
  } // namespace
  
  static coro::Task<void> childRequest(coro::Mutex& m, coro::Baton& b) {
    ShallowCopyRequestContextScopeGuard requestScope;
  
    auto* parentContext = dynamic_cast<TestRequestData*>(
        RequestContext::get()->getContextData(testToken1));
  
    EXPECT_TRUE(parentContext != nullptr);
  
    auto childKey = parentContext->key() + ".child";
  
    RequestContext::get()->setContextData(
        testToken2, std::make_unique<TestRequestData>(childKey));
  
    auto* childContext = dynamic_cast<TestRequestData*>(
        RequestContext::get()->getContextData(testToken2));
    CHECK(childContext != nullptr);
  
    {
      auto lock = co_await m.co_scoped_lock();
  
      CHECK_EQ(
          parentContext,
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData(testToken1)));
      CHECK_EQ(
          childContext,
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData(testToken2)));
  
      co_await b;
  
      CHECK_EQ(
          parentContext,
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData(testToken1)));
      CHECK_EQ(
          childContext,
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData(testToken2)));
    }
  
    CHECK_EQ(
        parentContext,
        dynamic_cast<TestRequestData*>(
            RequestContext::get()->getContextData(testToken1)));
    CHECK_EQ(
        childContext,
        dynamic_cast<TestRequestData*>(
            RequestContext::get()->getContextData(testToken2)));
  }
  
  static coro::Task<void> parentRequest(int id) {
    ShallowCopyRequestContextScopeGuard requestScope;
  
    // Should have captured the value at the time the coroutine was co_awaited
    // rather than at the time the coroutine was called.
    auto* globalData = dynamic_cast<TestRequestData*>(
        RequestContext::get()->getContextData("global"));
    CHECK(globalData != nullptr);
    CHECK_EQ("other value", globalData->key());
  
    std::string key = folly::to<std::string>("request", id);
  
    RequestContext::get()->setContextData(
        testToken1, std::make_unique<TestRequestData>(key));
  
    auto* contextData = RequestContext::get()->getContextData(testToken1);
    CHECK(contextData != nullptr);
  
    coro::Mutex mutex;
    coro::Baton baton1;
    coro::Baton baton2;
  
    auto fut1 = childRequest(mutex, baton1)
                    .scheduleOn(co_await coro::co_current_executor)
                    .start();
    auto fut2 = childRequest(mutex, baton1)
                    .scheduleOn(co_await coro::co_current_executor)
                    .start();
  
    CHECK_EQ(contextData, RequestContext::get()->getContextData(testToken1));
  
    baton1.post();
    baton2.post();
  
    (void)co_await std::move(fut1);
  
    CHECK_EQ(contextData, RequestContext::get()->getContextData(testToken1));
  
    // Check that context from child operation doesn't leak into this coroutine.
    CHECK(RequestContext::get()->getContextData(testToken2) == nullptr);
  
    (void)co_await std::move(fut2);
  
    // Check that context from child operation doesn't leak into this coroutine.
    CHECK(RequestContext::get()->getContextData(testToken2) == nullptr);
  }
  
  TEST(Task, RequestContextIsPreservedAcrossSuspendResume) {
    ManualExecutor executor;
  
    RequestContextScopeGuard requestScope;
  
    RequestContext::get()->setContextData(
        "global", std::make_unique<TestRequestData>("global value"));
  
    // Context should be captured at coroutine co_await time and not at
    // call time.
    auto task1 = parentRequest(1).scheduleOn(&executor);
    auto task2 = parentRequest(2).scheduleOn(&executor);
  
    {
      RequestContextScopeGuard nestedRequestScope;
  
      RequestContext::get()->setContextData(
          "global", std::make_unique<TestRequestData>("other value"));
  
      // Start execution of the tasks.
      auto fut1 = std::move(task1).start();
      auto fut2 = std::move(task2).start();
  
      // Check that the contexts set by starting the tasks don't bleed out
      // to the caller.
      CHECK(RequestContext::get()->getContextData(testToken1) == nullptr);
      CHECK(RequestContext::get()->getContextData(testToken2) == nullptr);
      CHECK_EQ(
          "other value",
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData("global"))
              ->key());
  
      executor.drain();
  
      CHECK(fut1.isReady());
      CHECK(fut2.isReady());
  
      // Check that the contexts set by the coroutines executing on the executor
      // do not leak out to the caller.
      CHECK(RequestContext::get()->getContextData(testToken1) == nullptr);
      CHECK(RequestContext::get()->getContextData(testToken2) == nullptr);
      CHECK_EQ(
          "other value",
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData("global"))
              ->key());
    }
  }
  
  TEST(Task, ContextPreservedAcrossMutexLock) {
    folly::coro::Mutex mutex;
  
    auto handleRequest =
        [&](folly::coro::Baton& event) -> folly::coro::Task<void> {
      RequestContextScopeGuard requestScope;
  
      RequestData* contextDataPtr = nullptr;
      {
        auto contextData = std::make_unique<TestRequestData>("some value");
        contextDataPtr = contextData.get();
        RequestContext::get()->setContextData(
            "mutex_test", std::move(contextData));
      }
  
      auto lock = co_await mutex.co_scoped_lock();
  
      // Check that the request context was preserved across mutex lock
      // acquisition.
      CHECK_EQ(
          RequestContext::get()->getContextData("mutex_test"), contextDataPtr);
  
      co_await event;
  
      // Check that request context was preserved across baton wait.
      CHECK_EQ(
          RequestContext::get()->getContextData("mutex_test"), contextDataPtr);
    };
  
    folly::ManualExecutor manualExecutor;
  
    folly::coro::Baton event1;
    folly::coro::Baton event2;
    auto t1 = handleRequest(event1).scheduleOn(&manualExecutor).start();
    auto t2 = handleRequest(event2).scheduleOn(&manualExecutor).start();
  
    manualExecutor.drain();
  
    event1.post();
  
    manualExecutor.drain();
  
    event2.post();
  
    manualExecutor.drain();
  
    EXPECT_TRUE(t1.isReady());
    EXPECT_TRUE(t2.isReady());
  
    EXPECT_FALSE(t1.hasException());
    EXPECT_FALSE(t2.hasException());
  }
  
  TEST(Task, RequestContextSideEffectsArePreserved) {
    auto f =
        [&](folly::coro::Baton& baton) -> folly::coro::detail::InlineTask<void> {
      RequestContext::create();
  
      RequestContext::get()->setContextData(
          testToken1, std::make_unique<TestRequestData>("test"));
  
      EXPECT_NE(RequestContext::get()->getContextData(testToken1), nullptr);
  
      // HACK: Need to use co_viaIfAsync() to ensure request context is preserved
      // across suspend-point.
      co_await folly::coro::co_viaIfAsync(
          &folly::InlineExecutor::instance(), baton);
  
      EXPECT_NE(RequestContext::get()->getContextData(testToken1), nullptr);
  
      co_return;
    };
  
    auto g = [&](folly::coro::Baton& baton) -> folly::coro::Task<void> {
      EXPECT_EQ(RequestContext::get()->getContextData(testToken1), nullptr);
      co_await f(baton);
      EXPECT_NE(RequestContext::get()->getContextData(testToken1), nullptr);
      EXPECT_EQ(
          dynamic_cast<TestRequestData*>(
              RequestContext::get()->getContextData(testToken1))
              ->key(),
          "test");
    };
  
    folly::ManualExecutor executor;
    folly::coro::Baton baton;
  
    auto t = g(baton).scheduleOn(&executor).start();
  
    executor.drain();
  
    baton.post();
  
    executor.drain();
  
    EXPECT_TRUE(t.isReady());
    EXPECT_FALSE(t.hasException());
  }
  
  TEST(Task, FutureTailCall) {
    EXPECT_EQ(
        42,
        folly::coro::blockingWait(
            folly::coro::co_invoke([&]() -> folly::coro::Task<int> {
              co_return co_await folly::makeSemiFuture().deferValue(
                  [](auto) { return folly::makeSemiFuture(42); });
            })));
  }
  
  // NOTE: This function is unused.
  // We just want to make sure this compiles without errors or warnings.
  folly::coro::Task<void>
  checkAwaitingFutureOfUnitDoesntWarnAboutDiscardedResult() {
    co_await folly::makeSemiFuture();
  
    using namespace std::literals::chrono_literals;
    co_await folly::futures::sleep(1ms);
  }
  
  folly::coro::Task<int&> returnIntRef(int& value) {
    co_return value;
  }
  
  TEST(Task, TaskOfLvalueReference) {
    int value = 123;
    auto&& result = folly::coro::blockingWait(returnIntRef(value));
    static_assert(std::is_same_v<decltype(result), int&>);
    CHECK_EQ(&value, &result);
  }
  
  TEST(Task, TaskOfLvalueReferenceAsTry) {
    folly::coro::blockingWait([]() -> folly::coro::Task<void> {
      int value = 123;
      auto&& result = co_await co_awaitTry(returnIntRef(value));
      CHECK(result.hasValue());
      CHECK_EQ(&value, &result.value().get());
  
      int& valueRef = co_await returnIntRef(value);
      CHECK_EQ(&value, &valueRef);
    }());
  }
  
  TEST(Task, CancellationPropagation) {
    folly::coro::blockingWait([]() -> folly::coro::Task<void> {
      auto token = co_await folly::coro::co_current_cancellation_token;
      CHECK(!token.canBeCancelled());
  
      folly::CancellationSource cancelSource;
  
      co_await folly::coro::co_withCancellation(
          cancelSource.getToken(), [&]() -> folly::coro::Task<void> {
            auto token2 = co_await folly::coro::co_current_cancellation_token;
            CHECK(token2.canBeCancelled());
            CHECK(!token2.isCancellationRequested());
  
            // The cancellation token should implicitly propagate into the
            //
            co_await[&]()->folly::coro::Task<void> {
              auto token3 = co_await folly::coro::co_current_cancellation_token;
              CHECK(token3 == token2);
              cancelSource.requestCancellation();
              CHECK(token3.isCancellationRequested());
            }
            ();
            CHECK(token2.isCancellationRequested());
          }());
    }());
  }
  
  TEST(Task, StartInlineUnsafe) {
    folly::coro::blockingWait([]() -> folly::coro::Task<void> {
      auto executor = co_await folly::coro::co_current_executor;
      bool hasStarted = false;
      bool hasFinished = false;
      auto sf = [&]() -> folly::coro::Task<void> {
        hasStarted = true;
        co_await folly::coro::co_reschedule_on_current_executor;
        hasFinished = true;
      }()
                             .scheduleOn(executor)
                             .startInlineUnsafe();
  
      // Check that the task started inline on the current thread.
      // It should not yet have completed, however, since the rest
      // of the coroutine needs this coroutine to suspend so the
      // executor can schedule the rest of it.
      CHECK(hasStarted);
      CHECK(!hasFinished);
  
      co_await std::move(sf);
  
      CHECK(hasFinished);
    }());
  }
  
  #endif
  