`exec::task` coroutine gets rescheduled when resuming

[kyusic]

Consider the following code using exec::task:

exec::task<void> h1() {
  std::cerr << "entered h1\n";
  co_return;
}

exec::task<void> h2() {
  co_await h1();
  std::cerr << "left h1\n";
}

When h1() completes and h2() resumes, I expected execution to continue on the same execution resource that was running h1(). However, it looks like exec::task reschedules onto its stored scheduler.
You can verify this by using a scheduler that prints a message when schedule() is called:

template <stdexec::scheduler BaseScheduler>
struct printing_scheduler {
  BaseScheduler base_sched;

  struct sender {
    using sender_concept = stdexec::sender_t;

    BaseScheduler base_sched;

    using base_sender_t = stdexec::schedule_result_t<BaseScheduler>;

    template <class Env>
    auto get_completion_signatures(Env &&) const noexcept
      -> stdexec::completion_signatures_of_t<base_sender_t, Env> {
      return {};
    }

    template <class Receiver>
    auto connect(Receiver &&rcvr) && {
      return stdexec::connect(
        stdexec::schedule(base_sched) | stdexec::then([] { std::cerr << "scheduled\n"; }),
        static_cast<Receiver &&>(rcvr));
    }

    constexpr auto get_env() const noexcept {
      return stdexec::env{
        stdexec::prop{
                      stdexec::get_completion_scheduler<stdexec::set_value_t>, printing_scheduler{base_sched}},
        stdexec::get_env(stdexec::schedule(base_sched))
      };
    }
  };

  auto schedule() const {
    return sender{base_sched};
  }

  bool operator==(const printing_scheduler &other) const = default;

  auto query(stdexec::get_forward_progress_guarantee_t q) const noexcept {
    return q(base_sched);
  }
};

static_assert(stdexec::scheduler<printing_scheduler<stdexec::inline_scheduler>>);
static_assert(stdexec::sender<printing_scheduler<stdexec::inline_scheduler>::sender>);

int main() {
  auto ctx = exec::single_thread_context{};
  stdexec::sync_wait(stdexec::starts_on(printing_scheduler{ctx.get_scheduler()}, h2()));
}

The output is:

scheduled
entered h1
scheduled   <- here you can see it gets rescheduled
left h1

I tried to track down the cause. When co_awaiting another exec::task inside an exec::task, it appears to run the following code:

stdexec/include/exec/task.hpp

Lines 392 to 400 in bb13042

	template <sender _Awaitable>
	requires __scheduler_provider<_Context>
	auto await_transform(_Awaitable&& __awaitable) noexcept -> decltype(auto) {
	// TODO: If we have a complete-where-it-starts query then we can optimize
	// this to avoid the reschedule
	return stdexec::as_awaitable(
	continues_on(static_cast<_Awaitable&&>(__awaitable), get_scheduler(*__context_)),
	*this);
	}

I suspect the rescheduling happens because continues_on() is used here.
Should the following overload be selected instead?

stdexec/include/exec/task.hpp

Lines 426 to 430 in bb13042

	template <class _Awaitable>
	auto await_transform(_Awaitable&& __awaitable) noexcept -> decltype(auto) {
	return with_awaitable_senders<__promise>::await_transform(
	static_cast<_Awaitable&&>(__awaitable));
	}

Thank you for your time.

[LeeHowes]

You write as if this is obviously wrong. Why do you expect it to continue on the same resource?

There are good reasons why rescheduling by default is the right choice. In normal code like:

void foo() {
// A
   bar();
// B
}

we don't exepect B to run on a different context from A, so we should assume that bar() does not change the execution resource. The default for coroutines should be similar: we should make sure that B runs on A's resource, and thus if bar ran elsewhere we should fix it. This is less efficient (often, at least), but a safer default.

In your example this is all on a single thread. So maybe your concern is simply that we are rescheduling when the execution resource is the same? It would be a valid optimization to not reschedule on a matching resource, but it wouldn't always be safe because we could end up with recursion depth issues or starvation given that a lot of use of coroutines is going to be on a shared resource where we want to add concurrency. If we never reschedule, we don't get concurrency.

exec::task is going for reasonable defaults with low risk of inexperienced users of coroutines encountering many of the issues that we've encountered with production coroutine deployments.

[kyusic]

Let’s look at the coroutine version of your example:

exec::task<void> foo() {
// A
    co_await bar();
// B
}

If bar() returns a general sender, I agree with you.
But what I mean here is the case where both bar() and foo() return an exec::task.

Let's say bar()'s implementation is again as follows;

exec::task<void> bar() {
// C
    co_await baz();
// D
}

Now, if baz() returns a sender(which may not be exec::task) with a different completion scheduler, rescheduling makes sure D runs on the same context as C — just like you said. All good there.
But if bar() itself doesn’t switch execution resources, then we can take advantage of that and skip the reschedule when resuming from bar() back to foo().

For reference, libunifex does the same optimization.

task::promise::await_transform

https://github.com/facebookexperimental/libunifex/blob/b6bedebc4d87eda5e31b364585a84576013eae67/include/unifex/task.hpp#L381-L385

      } else if constexpr (unifex::sender<Value>) {
        return unifex::await_transform(
            *this,
            with_scheduler_affinity(static_cast<Value&&>(value), this->sched_));

with_scheduler_affinity

https://github.com/facebookexperimental/libunifex/blob/b6bedebc4d87eda5e31b364585a84576013eae67/include/unifex/with_scheduler_affinity.hpp#L101-L113

struct _fn final {
  template(typename Sender, typename Scheduler)  //
      (requires sender<Sender> AND scheduler<Scheduler> AND
           tag_invocable<_fn, Sender, Scheduler>)  //
      constexpr auto
      operator()(Sender&& s, Scheduler&& sched) const
      noexcept(is_nothrow_tag_invocable_v<_fn, Sender, Scheduler>)
          -> tag_invoke_result_t<_fn, Sender, Scheduler> {
    // allow customization
    return tag_invoke(
        _fn{}, static_cast<Sender&&>(s), static_cast<Scheduler&&>(sched));
  }

task's customization of with_scheduler_affinity (scheduler is ignored)

https://github.com/facebookexperimental/libunifex/blob/b6bedebc4d87eda5e31b364585a84576013eae67/include/unifex/task.hpp#L906-L911

  template <typename Scheduler>
  friend typename _sa_task<T, nothrow>::type tag_invoke(
      tag_t<with_scheduler_affinity>, type&& task, Scheduler&&) noexcept {
    return {std::move(task)};
  }