tokio: distinguish LocalSet::enter() with being polled #6016

inq · 2023-09-19T11:26:13Z

Motivation

This fixes: #5020

Solution

If we enter a LocalSet, then the LocalEnterGuard sets CURRENT with it. But the CURRENT is being used with both of these two situations:

task::spawn_local()
being polled
But as we know, LocalEnterGuard should be related only with (1). When we schedule it, we should check if it is being polled, or just being entered. So I added a boolean variable to distinguish these situations.

hawkw

Overall, this change looks good to me --- thanks for adding a test reproducing the bug!

I had some minor suggestions for potentially simplifying some aspects of the implementation a bit. Let me know what you think!

tokio/tests/task_local_set.rs

hawkw · 2023-09-19T17:10:28Z

tokio/src/task/local.rs

+ let res = CURRENT.try_with(|LocalData { ctx, entered, .. }| {
 struct Reset<'a> {
 ctx_ref: &'a RcCell<Context>,
- val: Option<Rc<Context>>,
+ entered_ref: &'a Cell<bool>,
+ ctx_val: Option<Rc<Context>>,
+ entered_val: bool,
 }
 impl<'a> Drop for Reset<'a> {
 fn drop(&mut self) {
- self.ctx_ref.replace(self.val.take());
+ self.ctx_ref.set(self.ctx_val.take());
+ self.entered_ref.set(self.entered_val)
 }
 }
- let old = ctx.replace(Some(self.context.clone()));
+ let ctx_old = ctx.replace(Some(self.context.clone()));
+ let entered_old = entered.replace(false);

 let _reset = Reset {
 ctx_ref: ctx,
- val: old,
+ entered_ref: entered,
+ ctx_val: ctx_old,
+ entered_val: entered_old,


take it or leave it: the Reset code got substantially more complex in this change, and it's used in two different places. now that it's no longer fairly trivial, what do you think about moving this out of the with and with_if_possible functions, so that we can use one implementation in both places. Something like:

impl LocalData { fn enter(&self, context: Rc<Context>) -> Reset<'_> { let ctx_val = self.ctx.replace(Some(context)); let entered_val = self.entered.replace(false); Reset { ctx_ref: self.ctx, entered_ref: self.entered, ctx_val: ctx_old, entered_val: entered_old, } } } struct Reset<'a> { ctx_ref: &'a RcCell<Context>, entered_ref: &'a Cell<bool>, ctx_val: Option<Rc<Context>>, entered_val: bool, } impl<'a> Drop for Reset<'a> { fn drop(&mut self) { self.ctx_ref.set(self.ctx_val.take()); self.entered_ref.set(self.entered_val) } }

and then changing with and with_if_possible to just call

let _reset = local_data.enter(self.context.clone());

Thank you! I've moved out, and renamed it to more verbose name.

hawkw · 2023-09-19T17:11:24Z

tokio/src/task/local.rs

 ctx_ref: &'a RcCell<Context>,
- val: Option<Rc<Context>>,
+ entered_ref: &'a Cell<bool>,


nit, not a big deal: both ctx_ref and entered_ref reference fields of the LocalData. we could make this struct a word smaller by changing these to a single &'a LocalData reference, and just using that to access both ctx and entered.

It's a good idea. I've changed it!

hawkw · 2023-09-19T17:11:46Z

tokio/src/task/local.rs

-pub struct LocalEnterGuard(Option<Rc<Context>>);
+pub struct LocalEnterGuard {
+ ctx: Option<Rc<Context>>,
+ entered: bool,


can we add a comment here describing how entered is used and what it represents? thanks!

I've added! Thanks!

Sorry if this is overly nitpicky, but I think this comment could maybe be improved a bit --- it would be nice if it explained why we need to differentiate between enter and polling the LocalSet, and what behavior is controlled based on that. I don't think that's currelty all that clear from reading the code.

Oops, sorry! I've corrected it!

inq · 2023-09-19T18:52:31Z

@hawkw Thank you for providing an awesome review!

tokio/src/task/local.rs

hawkw · 2023-09-19T19:43:39Z

@hawkw Thank you for providing an awesome review!

@inq My pleasure, thanks for contributing the fix! :)

inq · 2023-09-25T01:51:40Z

@hawkw Hello! I am sorry but can we proceed this?

hawkw · 2023-09-28T17:58:41Z

@hawkw Hello! I am sorry but can we proceed this?

Hi, sorry for not getting back to you sooner --- I'm taking another look now!

hawkw

Overall, this looks good to me. I think it would be nice to document the behavior with the entered flag and when the LocalSet future is woken, but I feel good about the implementation and tests, and would be happy to merge this PR. Thank you!

hawkw · 2023-09-28T18:01:40Z

tokio/src/task/local.rs

-pub struct LocalEnterGuard(Option<Rc<Context>>);
+pub struct LocalEnterGuard {
+ ctx: Option<Rc<Context>>,
+ entered: bool,


Sorry if this is overly nitpicky, but I think this comment could maybe be improved a bit --- it would be nice if it explained why we need to differentiate between enter and polling the LocalSet, and what behavior is controlled based on that. I don't think that's currelty all that clear from reading the code.

tokio/src/task/local.rs

tokio/tests/task_local_set.rs

tokio/src/task/local.rs

Co-authored-by: Eliza Weisman <[email protected]>

inq · 2023-10-04T15:58:55Z

@hawkw I am really sorry for delayed works. Your kind correction helped me a lot! Thank you very much!

Darksonn · 2023-10-11T13:11:06Z

tokio/src/task/local.rs

+ /// Should be called except when we call `LocalSet::enter`.
+ /// Especially when we poll a LocalSet.
+ #[must_use = "dropping this guard will reset the entered state"]
+ fn enter(&self, ctx: Rc<Context>) -> LocalDataEnterGuard<'_> {
+ let ctx = self.ctx.replace(Some(ctx));
+ let entered = self.entered.replace(false);
+ LocalDataEnterGuard {
+ local_data_ref: self,
+ ctx,
+ entered,


I was reading over this code, and I'm confused about this. Why does calling enter set entered to false?

It's because of the name of the user-facing level functions. So this LocalData::enter should be called when we poll, and the value entered should be set only if user called the LocalSer::enter, which is public.
I agree that it's confused, but could not find suitable naming..

Can we rename the field to something like wake_on_spawn_local?

I agree with @Darksonn, the entered field could be renamed to something clearer.

Thank you very much! I've renamed it!
But... wake_on_spawn_local vs wake_on_schedule.. which would be better?

wake_on_schedule also makes sense to me.

Thank you! I've renamed it to wake_on_schdule!

Darksonn

Thanks.

Fix issue tokio-rs#5020

96bc15d

Darksonn added A-tokio Area: The main tokio crate M-task Module: tokio/task labels Sep 19, 2023

Merge branch 'master' into fix-local-enter-guard

b887953

inq changed the title ~~tokio: distinguish enter() with being polled~~ tokio: distinguish LocalSet::enter() with being polled Sep 19, 2023

inq added 2 commits September 19, 2023 22:29

Merge branch 'master' into fix-local-enter-guard

cdc5ca7

Merge branch 'master' into fix-local-enter-guard

8f86628

hawkw self-requested a review September 19, 2023 16:10

hawkw reviewed Sep 19, 2023

View reviewed changes

Fix

777efa0

Merge branch 'master' into fix-local-enter-guard

d2f542b

hawkw reviewed Sep 19, 2023

View reviewed changes

tokio/src/task/local.rs Outdated Show resolved Hide resolved

tokio/src/task/local.rs Show resolved Hide resolved

tokio/src/task/local.rs Show resolved Hide resolved

tokio/src/task/local.rs Show resolved Hide resolved

tokio/src/task/local.rs Outdated Show resolved Hide resolved

inq added 3 commits September 20, 2023 05:44

Update

c0d2634

Fmt

00f5b88

Split LocalDataEnterGuard with LocalEnterGuard

3cfc9f4

inq and others added 2 commits September 25, 2023 10:51

Merge branch 'master' into fix-local-enter-guard

ff24960

Merge branch 'master' into fix-local-enter-guard

95fa5ca

hawkw approved these changes Sep 28, 2023

View reviewed changes

inq and others added 3 commits October 4, 2023 22:52

Update tokio/tests/task_local_set.rs

5715a4c

Co-authored-by: Eliza Weisman <[email protected]>

Merge branch 'master' into fix-local-enter-guard

54657b9

Update

0616cbf

Darksonn reviewed Oct 11, 2023

View reviewed changes

inq added 3 commits October 16, 2023 01:09

Update

676318e

Merge branch 'master' into fix-local-enter-guard

869332f

wake_on_schedule

2af739b

Update

f61d1ed

Darksonn approved these changes Oct 15, 2023

View reviewed changes

Darksonn merged commit f3ad6cf into tokio-rs:master Oct 15, 2023
69 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

tokio: distinguish LocalSet::enter() with being polled #6016

tokio: distinguish LocalSet::enter() with being polled #6016

inq commented Sep 19, 2023 •

edited

Loading

hawkw left a comment

hawkw Sep 19, 2023

inq Sep 19, 2023

hawkw Sep 19, 2023

inq Sep 19, 2023 •

edited

Loading

hawkw Sep 19, 2023

inq Sep 19, 2023

hawkw Sep 28, 2023

inq Oct 4, 2023

inq commented Sep 19, 2023

hawkw commented Sep 19, 2023

inq commented Sep 25, 2023

hawkw commented Sep 28, 2023

hawkw left a comment

hawkw Sep 28, 2023

inq commented Oct 4, 2023

Darksonn Oct 11, 2023

inq Oct 12, 2023

Darksonn Oct 14, 2023

hawkw Oct 14, 2023

Nerdy5k Oct 15, 2023

inq Oct 15, 2023

Darksonn Oct 15, 2023

inq Oct 15, 2023

Darksonn left a comment

tokio: distinguish LocalSet::enter() with being polled #6016

tokio: distinguish LocalSet::enter() with being polled #6016

Conversation

inq commented Sep 19, 2023 • edited Loading

Motivation

Solution

hawkw left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inq Sep 19, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inq commented Sep 19, 2023

hawkw commented Sep 19, 2023

inq commented Sep 25, 2023

hawkw commented Sep 28, 2023

hawkw left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inq commented Oct 4, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Darksonn left a comment

Choose a reason for hiding this comment

inq commented Sep 19, 2023 •

edited

Loading

inq Sep 19, 2023 •

edited

Loading