HashMap iteration with retain while taking ownership

help
1

This is a simplification of the actual implementation, but I have a HashMap<String,(Arc<AtomicI8>, T)> that I iterate over frequently in a hot path. I check the atomic:

  • If it's 0, I keep the entry

  • If it's non-zero, I remove it, but I also want to take ownership of T to send to another thread for immediate processing

retain is an obvious choice, but it only gives out &mut (Arc<AtomicI8>, T), which is quite sad because, if the atomic is non-zero, I know I'll return false in the predicate and retain will drop the entry.

I tried using extract_if but benchmarks specific to my application showed higher variance and worse performance than retain (I suspect extract_if may be reshuffling the hash table internally?)

My current solution is to wrap the value in an Option

let map: HashMap<String, Option<(Arc<AtomicI8>, T)>> = HashMap::new();

// Values are always Some when inserted
// map.insert(key, Some((flag, data)));

map.retain(|key, entry| {
    if let Some(flag) = entry.as_ref().map(|(flag, _)| flag.load(Ordering::Relaxed)) {
        if flag == 0 {
            true
        } else {
            let (flag, data) = std::mem::take(entry).expect("It's always Some");
            process(data);
            false
        }
    } else {
        unreachable!()
    }
});

The values are always Some until mem::take sets them to None, after which retain removes the entry. The unreachable!() is there to satisfy the compiler, but it's unnecessary. So my question is there a better way to express this invariant to the compiler to avoid the Option wrapper and the unreachable!()? and is there a better extract and process during iteration pattern?

Would also appreciate any insight into what extract_if does under the hood and why it might be slower than retain.

It's my first post on this forum, please be nice!

2

This is the core implementation of extract_if:

I dont see why it would be slower than retain(). Removing an element from a hashmap is just writing a thombstone. It doesn't reshuffle anything.

3

I mean, from that snippet, extract_if looks like a lazy iterator that may or not be fully consumed. On the other hand retain simply knows what to keep or not decisively. There should be more instructions on extract_if branch compared to retain, from just that observation.

1 Like
4

Thank you, and sorry for the delayed reply

I think the key difference, if my understanding is correct, is that extract_if returns a lazy iterator that needs to be consumed for the elements to be actually removed from the original map, whereas retain does that eagerly.

5

it is true that extract_if is an iterator, while retain isn't.
extract_if must be an Iterator because it returns the element.
can you share how you did the comparison benches ?

i would personally try

extract_if(...).for_each(|(flag, data)| process(data))

as i heard it can be sometimes be nicer for perfs than a for loop, but it is a heuristic, so it might not help

6

that's exactly what I ended up doing, extract_if(...).for_each(|(flag, data)| process(data)), and I'm totally happy with the results.

In hindsight, my benchmarks were probably set up inaccurately. Basically, I was measuring the time difference between setting the atomic flag and calling process fn. If the resultant ExtractIf iterator isn't consumed right away, which it wasn't, this time difference will be large. So, in the context of my experiment, it wasn't exactly an apples-to-apples comparison between extract_if and retain