std::coroutine_handle<Promise>::done() returning unexpected value

326 Views Asked by At

I am trying to write a simple round-robin scheduler for coroutines. My simplified code is the following:

Generator<uint64_t> Counter(uint64_t i) {
  for (int j = 0; j < 2; j++) {
    co_yield i;
  }
}

...

void RunSimple(uint64_t num_coroutines) {
  std::vector<std::pair<uint64_t, Generator<uint64_t>>> gens;
  for (uint64_t i = 0; i < num_coroutines; i++) {
    // Storing coroutines into a vector is safe because we implemented the move
    // constructor and move assigment operator for Generator.
    gens.push_back({i, Counter(i)});
  }
  // This is a simple round-robin scheduler for coroutines.
  while (true) {
    for (auto it = gens.begin(); it != gens.end();) {
      uint64_t i = it->first;
      Generator<uint64_t>& gen = it->second;
      if (gen) {
        printf("Coroutine %lu: %lu.\n", i, gen());
        it++;
      } else {
        // `gen` has finished, so remove its pair from the vector.
        it = gens.erase(it);
      }
    }
    // Once all of the coroutines have finished, break.
    if (gens.empty()) {
      break;
    }
  }
}

When I run RunSimple(/*num_coroutines=*/4), I get the following output:

Coroutine 0: 0.
Coroutine 1: 1.
Coroutine 2: 2.
Coroutine 3: 3.
Coroutine 0: 0.
Coroutine 1: 1.
Coroutine 2: 2.
Coroutine 3: 3.
Coroutine 1: 1.
Coroutine 3: 3.
Coroutine 3: 3.

The last three lines of the output are unexpected... coroutines 1 and 3 do not seem to be exiting when I expect them to. Upon further investigation, this is happening because std::coroutine_handle<Promise>::done() is returning false for both of these coroutines. Do you know why this method is returning false... am I doing something wrong?


Edit: Here is my Generator implementation:

template <typename T>
struct Generator {
  struct promise_type;
  using handle_type = std::coroutine_handle<promise_type>;

  struct promise_type {
    T value_;
    std::exception_ptr exception_;

    Generator get_return_object() {
      return Generator(handle_type::from_promise(*this));
    }
    std::suspend_always initial_suspend() { return {}; }
    std::suspend_always final_suspend() noexcept { return {}; }
    void unhandled_exception() { exception_ = std::current_exception(); }
    template <std::convertible_to<T> From>  // C++20 concept
    std::suspend_always yield_value(From&& from) {
      value_ = std::forward<From>(from);
      return {};
    }
    void return_void() {}
  };

  handle_type h_;

  Generator(handle_type h) : h_(h) {
  }
  Generator(Generator&& g) : h_(std::move(g.h_)) { g.h_ = nullptr; }
  ~Generator() {
    if (h_) {
      h_.destroy();
    }
  }

  Generator& operator=(Generator&& g) {
    h_ = std::move(g.h_);
    g.h_ = nullptr;
    return *this;
  }
  explicit operator bool() {
    fill();
    return !h_.done();
  }
  T operator()() {
    fill();
    full_ = false;
    return std::move(h_.promise().value_);
  }

 private:
  bool full_ = false;

  void fill() {
    if (!full_) {
      h_();
      if (h_.promise().exception_)
        std::rethrow_exception(h_.promise().exception_);
      full_ = true;
    }
  }
};
1

There are 1 best solutions below

5
On BEST ANSWER

Your program contains undefined behavior.

The issue is your fill resumes the coroutine when !full_, regardless of whether the coroutine is at final suspend or not, which you can learn by calling h_.done().

Generator<uint64_t> Counter(uint64_t i) {
  for (int j = 0; j < 2; j++) {
    co_yield i;
  }
  // final suspend
}

If you resume the coroutine at the final suspend point, it will destory itself and you can no longer do anything with the coroutine handle you have.

And you call fill from operator bool, meaning that it, when called while the coroutine is suspended at final suspend, first destroys the coroutine, and then tries to access it, which is UB:

    fill(); // potentially destroy the coroutine
    return !h_.done(); // access the destroyed coroutine

You could fix this by making fill aware of doneness:

   void fill() {
    if (!h_.done() && !full_) {
      h_();
      if (h_.promise().exception_)
        std::rethrow_exception(h_.promise().exception_);
      full_ = true;
    }
  }

Also, your move assignment operator leaks the current coroutine:

  Generator& operator=(Generator&& g) {
    h_ = std::move(g.h_); // previous h_ is leaked
    g.h_ = nullptr;
    return *this;
  }

You probably want to have something like if (h_) h_.destroy(); in the beginning.

Also, as mentioned in the comments, the full_ member has to be carried over in the move constructor and assignment operators:

Generator(Generator&& g) 
  : h_(std::exchange(g.h_, nullptr))
  , full_(std::exchange(g.full_, false)) {}

Generator& operator=(Generator&& g) {
  full_ = std::exchange(g.full_, false);
  ...
}