Skip to content

[k2] add await set primitive#1418

Merged
astrophysik merged 17 commits into
masterfrom
vsadokhov/k2-await-set
Oct 20, 2025
Merged

[k2] add await set primitive#1418
astrophysik merged 17 commits into
masterfrom
vsadokhov/k2-await-set

Conversation

@astrophysik
Copy link
Copy Markdown
Contributor

@astrophysik astrophysik commented Oct 3, 2025

General

This pull request adds an equivalent of JoinSet from rust-tokio to runtime-light. This primitive allows coroutines to be dynamically added to it and waits for their completion.

Example:

kphp::coro::await_set<int64_t> await_set{};
for (int i = 0; i < 5; ++i) {
    await_set.push(coroutine_call());
}

while (!await_set.empty()) {
    auto res = co_await await_set.next();
    // process result here
}

The primitive supports multiple awaiters. All awaiters are stored in a linked list.

In the constructor, you can also set the waiting policy. Two options are available - suspend_on_empty and resume_on_empty.
They affect two aspects:

  1. What happens when waiting begins on an empty set
  2. What happens when the set becomes empty during waiting.

Details

The implementation of the primitive is similar to the implementations of when_any and when_all. The awaitables passed to push are wrapped in an await_set_task, in the final_suspend of which the result is stored. Awaiters are linked to the list of results through the await_broker class, which maintains a singly linked list of completed tasks and a doubly linked list of awaiters. As a remark, I would like to note that awaiters and ready tasks are stored in the lists as LIFO.

Motivation

This is the first step towards implementing wait_queue builtins in the new runtime. Here's how a wait_queue can be approximately implemented using a primitive.

inline kphp::coro::shared_task<int64_t> fork_function() {  
  co_await kphp::coro::io_scheduler::get().yield_for(10ms);  
  co_return 42;  
}  
  
  
inline kphp::coro::task<kphp::coro::shared_task<>> helper_task(kphp::coro::shared_task<int64_t> shared_task) {  
  co_await shared_task.when_ready();  
  co_return shared_task;  
}

auto shared_task = fork_function();
kphp::coro::await_set<kphp::coro::shared_task<>> wait_queue;
wait_queue.push(helper_task(shared_task))
auto erased_task = co_await wait_queue.next();
auto task = static_cast<kphp::coro::shared_task<int64_t>>(std::move(erased_task));  
auto res = co_await task;  
f$var_dump(res);

@astrophysik astrophysik requested a review from apolyakov October 3, 2025 16:52
@astrophysik astrophysik self-assigned this Oct 3, 2025
@astrophysik astrophysik added the k2 Affects compiler or runtime in K2 mode label Oct 3, 2025
@astrophysik astrophysik force-pushed the vsadokhov/k2-await-set branch from 4520095 to 8c3a28b Compare October 9, 2025 11:35
@astrophysik astrophysik marked this pull request as ready for review October 9, 2025 12:06
@astrophysik astrophysik added the enhancement New feature or request label Oct 14, 2025
Comment thread runtime-light/coroutine/await-set.h
Comment thread runtime-light/coroutine/await-set.h Outdated
Comment thread runtime-light/coroutine/await-set.h Outdated
Comment thread runtime-light/coroutine/await-set.h
Comment thread runtime-light/coroutine/await-set.h
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
@astrophysik astrophysik requested a review from apolyakov October 15, 2025 09:18
@apolyakov
Copy link
Copy Markdown
Contributor

Summary after discussion:

  1. There should be a single policy of how await_set works
  2. abort_all should not prevent awaiters from awaiting on the same await_set
  3. await_set should be movable
  4. kphp::coro should only contain await_set, no other symbols are required in public API

Comment thread runtime-light/coroutine/await-set.h Outdated
Comment thread runtime-light/coroutine/await-set.h Outdated
Comment thread runtime-light/coroutine/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h Outdated
Comment thread runtime-light/coroutine/detail/await-set.h
Copy link
Copy Markdown
Contributor

@apolyakov apolyakov left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@astrophysik astrophysik added this to the next milestone Oct 20, 2025
@astrophysik astrophysik merged commit 48be0c5 into master Oct 20, 2025
5 of 7 checks passed
@astrophysik astrophysik deleted the vsadokhov/k2-await-set branch October 20, 2025 15:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request k2 Affects compiler or runtime in K2 mode

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants