Remove `WaitMap` dependency #1183

charliermarsh · 2024-01-30T05:02:22Z

Summary

This is an attempt to #1163 by removing the WaitMap and gaining more granular control over the values that we hold over await boundaries.

charliermarsh · 2024-01-30T05:02:38Z

@konstin - I haven't done any benchmarking, but this does consistently not-deadlock for me.

crates/once-map/src/lib.rs

konstin · 2024-01-30T11:37:57Z

Current dependencies on/for this PR:

main
- PR Remove WaitMap dependency #1183 👈
  - PR Yield after channel send and move cpu tasks to thread #1163

This stack of pull requests is managed by Graphite.

konstin · 2024-01-30T11:39:20Z

(Sorry for pushing this branch, i misread the updated graphite docs)

konstin

This removes the deadlock! Benchmarks in the upstack PR are looking good. My main concern is deadlock-warning in the DashMap::entry api

konstin · 2024-01-30T11:34:22Z

crates/once-map/src/lib.rs

+            Value::Waiting(notify) => {
+                let notify = notify.clone();
+                drop(entry);
+                notify.notified().await;


Can we assert that on drop, nobody is waiting anymore?

Can you say more?

If there is a task waiting on some request but there is no task providing it (because we're done, we're dropping the once map), this sounds like a bug.

Oh! I see now. Yes. On Drop of the OnceMap.

I agree it might be a bug... but I could also see that maybe it isn't? Maybe the caller has finished what they needed and didn't need to wait for everything to finish. And it is perhaps hard to distinguish between "there is no task providing it" and "there is a task providing it, but it hasn't done so yet." I defer to y'all here because I don't have enough context on how this is used. But I absolutely agree that if you can assert that any Waiting values are a bug, then it might be nice to assert it on Drop. (Although note that a panic during Drop is an instant abort.)

konstin · 2024-01-30T11:35:53Z

crates/puffin-resolver/src/resolver/mod.rs

@@ -406,6 +405,9 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                if self.index.distributions.register_owned(dist.package_id()) {
                    priorities.add(dist.name().clone());
                    request_sink.unbounded_send(Request::Dist(dist))?;
+
+                    // Yield, to allow subscribers to continue.


I'd add here that we specifically need this because the channel is sync

konstin · 2024-01-30T11:51:51Z

crates/once-map/src/lib.rs

-        let mut lock = self.started.lock().unwrap();
-        if lock.contains(key) {
-            return false;
+        let entry = self.items.entry(key.to_owned());


The docs for entry say:

Locking behaviour: May deadlock if called when holding any sort of reference into the map.

This assumes we never actually call this function from two threads at the same time. This is the same call that WaitMap previously deadlocked in.

If we know that the map will stay on the main thread anyway, we can use Cell<FxHashMap<K, V>> instead of DashMap. If we want share the index between threads (i don't know that - it depends on how we want to structure our async code) we'll need something like RwLock<FxHashMap<K, V>>. DashMap is a Box<[RwLock<HashMap<K, V, S>>]> internally (https://docs.rs/dashmap/latest/src/dashmap/lib.rs.html#88-92) and our lock times are minimal, so i expect no perf difference. DashMap is missing a get_or_insert() -> bool, an "atomic" compare-and-swap option, that both we and WaitMap would need to be correct.

I don't think the warning is about calling entry from two different threads simultaneously. That's normal and should be fine, otherwise it'd be a pretty poor concurrent hashmap. The warning is, AIUI, about calling entry (or get) while you have a reference to a previous call in hand in the same thread. When multiple threads are calling it, one will (hopefully) eventually make progress, drop the reference and unblock the other thread. But if you do let entry = self.items.entry(foo); and then self.items.get(foo) while entry is still alive, then that get call seems likely to block waiting for entry to drop. Which, of course, will never happen because the thread is blocked on get.

Thanks for clarifying, this makes much more sense! I was indeed assuming that dashmap was behaving rather poorly, but this makes more sense and works for us if we don't have any await points while we hold the entry.

Thinking about this longer, was the problem with waitmap maybe merely that we there was some yielding while holding an entry?

@konstin - I spent a while investigating that, and trying to make changes to the resolver to solve it, but I ultimately couldn't figure out where it might be.

I also briefly looked at waitmap's implementation and nothing jumped out at me immediately.

konstin · 2024-01-30T11:55:28Z

crates/once-map/src/lib.rs

-        self.wait_map.insert(key, value);
+        if let Some(Value::Waiting(notify)) = self.items.insert(key, Value::Filled(Arc::new(value)))
+        {
+            notify.notify_waiters();


Should we also yield here? Notifying is surprisingly sync, so the subscribers won't get notified immediately until the next yielding of the task that called done. (This would make the function async, but i think it's correct that all operations on the OnceMap should be async)

Notifying is surprisingly sync, so the subscribers won't get notified immediately until the next yielding of the task that called done.

Hmmm. Are you sure? We're using the multi-threaded runtime for tokio right? If so, AIUI, other waiters could be notified and acting on it before notify.notify_waiters() even finishes.

Yep that's correct, but our receivers are all on the some thread atm as far as i can see.

Is there a reason the waiters need to wake up sooner?

In #1163 i got a speedup from 30ms to 20ms for warm cache jupyter by inserting one tokio::task::yield_now().await, now i'm motivated to avoid these kinds of bottlenecks.

BurntSushi

This is great. I like the use of Arc to simplify things here.

BurntSushi · 2024-01-30T13:11:46Z

crates/once-map/src/lib.rs

-        let mut lock = self.started.lock().unwrap();
-        if lock.contains(key) {
-            return false;
+        let entry = self.items.entry(key.to_owned());


I don't think the warning is about calling entry from two different threads simultaneously. That's normal and should be fine, otherwise it'd be a pretty poor concurrent hashmap. The warning is, AIUI, about calling entry (or get) while you have a reference to a previous call in hand in the same thread. When multiple threads are calling it, one will (hopefully) eventually make progress, drop the reference and unblock the other thread. But if you do let entry = self.items.entry(foo); and then self.items.get(foo) while entry is still alive, then that get call seems likely to block waiting for entry to drop. Which, of course, will never happen because the thread is blocked on get.

crates/once-map/src/lib.rs

BurntSushi · 2024-01-30T13:21:37Z

crates/once-map/src/lib.rs

-        self.wait_map.insert(key, value);
+        if let Some(Value::Waiting(notify)) = self.items.insert(key, Value::Filled(Arc::new(value)))
+        {
+            notify.notify_waiters();


Notifying is surprisingly sync, so the subscribers won't get notified immediately until the next yielding of the task that called done.

Hmmm. Are you sure? We're using the multi-threaded runtime for tokio right? If so, AIUI, other waiters could be notified and acting on it before notify.notify_waiters() even finishes.

crates/once-map/src/lib.rs

BurntSushi · 2024-01-30T13:23:58Z

crates/once-map/src/lib.rs

+            Value::Waiting(notify) => {
+                let notify = notify.clone();
+                drop(entry);
+                notify.notified().await;


Can you say more?

crates/once-map/src/lib.rs

charliermarsh · 2024-01-30T13:50:13Z

Thank you both so much for the close read, I needed it!

zanieb · 2024-01-30T15:21:41Z

crates/once-map/src/lib.rs

+            dashmap::mapref::entry::Entry::Occupied(_) => false,
+            dashmap::mapref::entry::Entry::Vacant(entry) => {
+                entry.insert(Value::Waiting(Arc::new(Notify::new())));
+                true


This could return some Handle object that we could use to enforce the invariant, right? e.g. done would consume a handle which would have a reference to the key instead and we'd be able to check if the handle was not used on drop.

If this method returns true, you need to start a job and call [OnceMap::done] eventually
or other tasks will hang.

Not sure how problematic that is.

I agree that would be a much nicer API. Will consider it in a future PR, the dataflow might be tedious for now.

This reverts commit ea97c30185d08f8ce736b021f54e95859d54b086.

charliermarsh requested a review from konstin January 30, 2024 05:02

charliermarsh commented Jan 30, 2024

View reviewed changes

crates/once-map/src/lib.rs Show resolved Hide resolved

konstin force-pushed the charlie/wait branch from 71579fa to 2ab1949 Compare January 30, 2024 11:37

konstin mentioned this pull request Jan 30, 2024

Yield after channel send and move cpu tasks to thread #1163

Merged

konstin reviewed Jan 30, 2024

View reviewed changes

BurntSushi reviewed Jan 30, 2024

View reviewed changes

zanieb reviewed Jan 30, 2024

View reviewed changes

konstin approved these changes Jan 30, 2024

View reviewed changes

charliermarsh added 8 commits January 30, 2024 14:54

Add yields

fdea295

Try to change once-map

a54943c

Revert "Try to change once-map"

965e835

This reverts commit ea97c30185d08f8ce736b021f54e95859d54b086.

Don't hold across awaits

e66c6ea

Try to add a waiting map

e20ef9f

This seems to work

7929341

Finish impl

d1408cf

Address review feedback

887fc3f

charliermarsh force-pushed the charlie/wait branch from 2ab1949 to 887fc3f Compare January 30, 2024 20:11

charliermarsh merged commit 3f5e730 into main Jan 30, 2024
5 checks passed

charliermarsh deleted the charlie/wait branch January 30, 2024 20:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Remove `WaitMap` dependency #1183

Remove `WaitMap` dependency #1183

charliermarsh commented Jan 30, 2024

charliermarsh commented Jan 30, 2024

konstin commented Jan 30, 2024 •

edited

Loading

konstin commented Jan 30, 2024

konstin left a comment

konstin Jan 30, 2024

BurntSushi Jan 30, 2024

konstin Jan 30, 2024

BurntSushi Jan 30, 2024

konstin Jan 30, 2024

konstin Jan 30, 2024

BurntSushi Jan 30, 2024

konstin Jan 30, 2024

konstin Jan 30, 2024

charliermarsh Jan 30, 2024

BurntSushi Jan 30, 2024

konstin Jan 30, 2024

BurntSushi Jan 30, 2024

konstin Jan 30, 2024 •

edited

Loading

zanieb Jan 30, 2024

konstin Jan 30, 2024

BurntSushi left a comment

BurntSushi Jan 30, 2024

BurntSushi Jan 30, 2024

BurntSushi Jan 30, 2024

charliermarsh commented Jan 30, 2024

zanieb Jan 30, 2024 •

edited

Loading

charliermarsh Jan 30, 2024

Remove WaitMap dependency #1183

Remove WaitMap dependency #1183

Conversation

charliermarsh commented Jan 30, 2024

Summary

charliermarsh commented Jan 30, 2024

konstin commented Jan 30, 2024 • edited Loading

konstin commented Jan 30, 2024

konstin 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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

konstin Jan 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BurntSushi 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

charliermarsh commented Jan 30, 2024

zanieb Jan 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Remove `WaitMap` dependency #1183

Remove `WaitMap` dependency #1183

konstin commented Jan 30, 2024 •

edited

Loading

konstin Jan 30, 2024 •

edited

Loading

zanieb Jan 30, 2024 •

edited

Loading