Where the tokio task context switch code?

i am trying to read the source code of tokio. for. i want to get a deeper understand of rust concurrency model, but i can't find the source code of tokio's task context switch.

here some code i had read:

  1. tokio::runtime::Schedule
  2. the way tokio bind a future to a Task
  3. taks's vtable

and i find the task.poll() finally enter this method poll and then enter poll_inner

from the callstack i can't find any context switch when os thread A from run Task a switch to run Task b.

So my question is:
1. does the callstack i found is right?
2. where is the code of tokio's Task context switch?
3. where the locationTask's context data store at (as far as i know, the information for task to be run able such as the computer register's data should store at process's heap location)

Thanks for your reply. Thanks

In rust async functions don't use green threads with stack switching. Instead async functions are codegened to state machines which advance to the next state every time you call poll on them. All local variables that live across .await are stored directly in the value implementing Future and all local variables that don't live across .await are stored on the regular stack which is popped before returning from .poll().

6 Likes

The only thing you need to do to context switch is call poll on the next future that should execute.

2 Likes

So, what's you meaning is: a Task is just a function/struct pointer, when an os thread execute a task , it. just run this function(which jump to the function's pc address)

Do i got what you mean?

Thanks

I'm sure you're aware, but just to clarify this point, this is something of a philosophical/semantic issue - either way you end up with a stack of local variables once you unwind all the types.

For example:

async fn foo() {
  let local = 123;
  bar().await;
  println!("{local}");
}

async fn bar() {
  let other = 456;
  yield().await;
  println!("{other}");
}

In the above, the generated return types are similar to:

enum FooState {
  Initial,
  AwaitBar { local: i32, bar: BarState },
}

enum BarState {
  Initial,
  AwaitYield { other: i32, yield: YeildState },
}

...

so the result of foo() after a poll(), where it has hit yield(), should be:

FooState::AwaitBar {
  local: 123,
  bar: BarState::AwaitYield {
    other: 345,
    yield: ...
  },
}

This is directly equivalent to a callstack like:

yield
bar { other: 345 }
foo { local: 123 }

the main difference being the pointer is the outermost frame, rather than the innermost (presumably inlining takes care of a lot of that).

so, if a os thread A, run this future, Does a struct FooState::AwaitBar will allocate at global heap?
and if this is right, does a future have so many nesting .await call will make a big memory allocate ?

No, the Future implementation that async fn returns directly contains all the state, so it will end up on the stack by default.

This is why you can't directly recurse async fns:

async fn fac(n: i32) -> i32 {
  if n <= 1 {
    1
  } else {
    n * fac(n - 1).await
  }
}

gives:

error[E0733]: recursion in an `async fn` requires boxing
 --> src/lib.rs:1:25
  |
1 | async fn fac(n: i32) -> i32 {
  |                         ^^^ recursive `async fn`
  |
  = note: a recursive `async fn` must be rewritten to return a boxed `dyn Future`
  = note: consider using the `async_recursion` crate: https://crates.io/crates/async_recursion

The "rewritten to return a boxed dyn Future" code looks like:

use std::future::Future;
use std::pin::Pin;

fn fac(n: i32) -> Pin<Box<dyn Future<Output = i32>>> {
  return Box::pin(imp(n)); // puts the implicit future struct on the stack
  
  async fn imp(n: i32) -> i32 {
      if n <= 1 {
        1
      } else {
        n * fac(n - 1).await
      }
  }
}

which is (presumably) what the async_recursion crate does.

If you do end up with very deeply nested async fn calls, it will end up with large state, but not really any more than would end up on a sync stack. You can simply do something like the above with a Box::pin(big_async_fn_state()).await to break it up, if you want to.

thank you very much, thanks for your such patient explanation.

if a Future store at stack, i think it can't pass between os thread, for i can't imagine a function run on. different os thread without keep the cpu context the same. it's right?

if. previous conclusion is right. if we want task can be passed at different os thread, it should store at heap ?

And i also have some confusing things. about Pin, i read this async-book and i got the exist meaning of Pin, but. i can't find where the use of mem::swap.

Thanks for your replay

It's not that they are stored on the stack, but that they store their local state inline in the value returned by the function. So when you pass that to spawn, that state is moved into tokio's runtime machinery, which will presumably box the argument (though perhaps not directly) so it can handle lots of different future types being spawned. When building for release, all this should end up inlined and the state will be directly initialized on the heap.

The end result is that you only have to box and indirectly poll a future when it's actually necessary, minimizing the amount of heap allocation overhead and pointer jumping that needs to happen when executing async code.

So, to be clear, by default spawn will cause the future to be picked up by any of tokio's worker threads, whenever one becomes available. Sometimes that's not safe, and you'll get an unfortunately confusing error about your future not implementing Send. Fixing that depends on the context, it generally means you're doing something like trying to hold a Mutex across an .await.

The Pin API is infamously confusing. I would mostly treat it as an internal detail necessary for await to work: but it basically boils down to the fact that futures are allowed to hold self references, because you want to be able to have a local variable referencing another inside async functions. To do this, futures require that you promise (no pun intended) that you won't move them once you start executing them, and that promise looks like putting the future in a Pin type. Most of the time that looks like std::pin::pin!() if the future will not be returned or otherwise moved from the function it's in, or Box::pin() if it will. It gets a lot messier if you're implementing a Future directly.

1 Like

Thanks again. your replies help me to build a full image of rust async.

any advise for a more deeper understand of rust async.

If you want some real gory details about how to use async and what can go wrong, Jon Gjengset's Crust of Rust stream on async is, as usual for him, excellent: Crust of Rust: async/await - YouTube

And if you want to go on an adventure, Amos is always fun, I suggest starting with the aptly named Pin and suffering. Pin and suffering

These two are where at least 70% of my Rust knowledge come from, so I really recommend exploring pretty much everything they've put out, even if they create quite lengthy content.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.