Skip to content

Commit 4889717

Browse files
maciejbacalextvinodreddy-g
maciejbacalext
authored and
vinodreddy-g
committed
orch: Make the select macro work with async functions and closuers
1 parent 31ca7c5 commit 4889717

1 file changed

Lines changed: 92 additions & 28 deletions

File tree

‎src/async_runtime/src/macros/select.rs‎

Lines changed: 92 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -115,10 +115,16 @@ macro_rules! select {
115115
// Main logic. This is matched at the end of the recursion.
116116
(@ ($($max_index_underscores:tt)*); $(($($index_underscores:tt)*) $v:pat = $f:expr => $b:block,)+) => {{
117117
use core::future::Future;
118+
use core::pin::pin;
118119

119120
// If a bit at case's index is set to 1, the case is disabled.
120121
let mut disabled_cases = 0_usize;
121122

123+
// The pin! macro is necessary to pin futures that don't implement that Unpin trait, for example async functions and closures.
124+
// The use of the pin! macro is safe because select! awaits on the poll_fn immediately instad of returning the future.
125+
// If select was to return the future, the wrapping block would need to be async to capture the local variable crated by pin!.
126+
let mut pins = ($(pin!($f),)+);
127+
122128
core::future::poll_fn(|context| {
123129
// Create a block for each case and poll its future.
124130
// TODO: This will have to be a for loop over all futures if randomization is needed.
@@ -127,7 +133,7 @@ macro_rules! select {
127133
let case_mask = 1 << case_index;
128134

129135
if (disabled_cases & case_mask) == 0 {
130-
let mut pin = core::pin::Pin::new(&mut $f);
136+
let ($($index_underscores,)* pin, ..) = &mut pins;
131137

132138
match pin.as_mut().poll(context) {
133139
core::task::Poll::Ready(result) => {
@@ -360,7 +366,7 @@ mod tests {
360366
#[test]
361367
fn one_case() {
362368
async fn test() -> i32 {
363-
let mut fut1 = future::ready(111);
369+
let fut1 = future::ready(111);
364370
let mut result = 0;
365371

366372
select! {
@@ -381,9 +387,9 @@ mod tests {
381387
#[test]
382388
fn case_1_first() {
383389
async fn test() -> i32 {
384-
let mut fut1 = future::ready(111);
385-
let mut fut2 = future::pending::<i32>();
386-
let mut fut3 = future::pending::<i32>();
390+
let fut1 = future::ready(111);
391+
let fut2 = future::pending::<i32>();
392+
let fut3 = future::pending::<i32>();
387393
let mut result = 0;
388394

389395
select! {
@@ -410,9 +416,9 @@ mod tests {
410416
#[test]
411417
fn case_2_first() {
412418
async fn test() -> i32 {
413-
let mut fut1 = future::pending::<i32>();
414-
let mut fut2 = future::ready(222);
415-
let mut fut3 = future::pending::<i32>();
419+
let fut1 = future::pending::<i32>();
420+
let fut2 = future::ready(222);
421+
let fut3 = future::pending::<i32>();
416422
let mut result = 0;
417423

418424
select! {
@@ -439,9 +445,9 @@ mod tests {
439445
#[test]
440446
fn case_3_first() {
441447
async fn test() -> i32 {
442-
let mut fut1 = future::pending::<i32>();
443-
let mut fut2 = future::pending::<i32>();
444-
let mut fut3 = future::ready(333);
448+
let fut1 = future::pending::<i32>();
449+
let fut2 = future::pending::<i32>();
450+
let fut3 = future::ready(333);
445451
let mut result = 0;
446452

447453
select! {
@@ -468,10 +474,10 @@ mod tests {
468474
#[test]
469475
fn case_2_and_3_ready() {
470476
async fn test() -> i32 {
471-
let mut fut1 = future::pending::<i32>();
477+
let fut1 = future::pending::<i32>();
472478
// Both futures are ready, but since there's no randomization, the 2nd will return first.
473-
let mut fut2 = future::ready(222);
474-
let mut fut3 = future::ready(333);
479+
let fut2 = future::ready(222);
480+
let fut3 = future::ready(333);
475481
let mut result = 0;
476482

477483
select! {
@@ -498,10 +504,10 @@ mod tests {
498504
#[test]
499505
fn optional_case_matches() {
500506
async fn test() -> i32 {
501-
let mut fut1 = future::pending::<i32>();
502-
let mut fut2: future::Ready<Option<i32>> = future::ready(Some(222));
503-
let mut fut3 = future::ready(333);
504-
let mut fut4 = future::pending::<i32>();
507+
let fut1 = future::pending::<i32>();
508+
let fut2: future::Ready<Option<i32>> = future::ready(Some(222));
509+
let fut3 = future::ready(333);
510+
let fut4 = future::pending::<i32>();
505511
let mut result = 0;
506512

507513
select! {
@@ -531,10 +537,10 @@ mod tests {
531537
#[test]
532538
fn optional_case_fails_to_match() {
533539
async fn test() -> i32 {
534-
let mut fut1 = future::pending::<i32>();
535-
let mut fut2: future::Ready<Option<i32>> = future::ready(None);
536-
let mut fut3 = future::ready(333);
537-
let mut fut4 = future::pending::<i32>();
540+
let fut1 = future::pending::<i32>();
541+
let fut2: future::Ready<Option<i32>> = future::ready(None);
542+
let fut3 = future::ready(333);
543+
let fut4 = future::pending::<i32>();
538544
let mut result = 0;
539545

540546
select! {
@@ -564,9 +570,9 @@ mod tests {
564570
#[test]
565571
fn all_cases_pending() {
566572
async fn test() -> i32 {
567-
let mut fut1 = future::pending::<i32>();
568-
let mut fut2 = future::pending::<i32>();
569-
let mut fut3 = future::pending::<i32>();
573+
let fut1 = future::pending::<i32>();
574+
let fut2 = future::pending::<i32>();
575+
let fut3 = future::pending::<i32>();
570576
let mut result = 0;
571577

572578
select! {
@@ -596,9 +602,67 @@ mod tests {
596602
#[test]
597603
fn select_returns_result() {
598604
async fn test() -> i32 {
599-
let mut fut1 = future::ready(111);
600-
let mut fut2 = future::pending::<i32>();
601-
let mut fut3 = future::pending::<i32>();
605+
let fut1 = future::ready(111);
606+
let fut2 = future::pending::<i32>();
607+
let fut3 = future::pending::<i32>();
608+
609+
select! {
610+
var1 = fut1 => {
611+
var1
612+
}
613+
_ = fut2 => {
614+
222
615+
}
616+
_ = fut3 => {
617+
333
618+
}
619+
}
620+
}
621+
622+
let mut context = task::Context::from_waker(task::Waker::noop());
623+
let mut future = Box::pin(test());
624+
let result = future.as_mut().poll(&mut context);
625+
assert_eq!(result, task::Poll::Ready(111));
626+
}
627+
628+
#[test]
629+
fn select_called_with_async_function() {
630+
async fn test() -> i32 {
631+
async fn first() -> i32 {
632+
111
633+
}
634+
635+
let fut1 = first();
636+
let fut2 = future::pending::<i32>();
637+
let fut3 = future::pending::<i32>();
638+
639+
select! {
640+
var1 = fut1 => {
641+
var1
642+
}
643+
_ = fut2 => {
644+
222
645+
}
646+
_ = fut3 => {
647+
333
648+
}
649+
}
650+
}
651+
652+
let mut context = task::Context::from_waker(task::Waker::noop());
653+
let mut future = Box::pin(test());
654+
let result = future.as_mut().poll(&mut context);
655+
assert_eq!(result, task::Poll::Ready(111));
656+
}
657+
658+
#[test]
659+
fn select_called_with_async_closure() {
660+
async fn test() -> i32 {
661+
let first = async || 111;
662+
663+
let fut1 = first();
664+
let fut2 = future::pending::<i32>();
665+
let fut3 = future::pending::<i32>();
602666

603667
select! {
604668
var1 = fut1 => {

0 commit comments

Comments
 (0)