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!

1 post - 1 participant

Read full topic