Concurrency: Designing a TaintableMap

[kianenigma]

I have a particular data structure in mind, which resonates a bit with a concurrent HashMap, yet it is not quite the same.

What I want is a map, where each first read of a key will taint that key with some identifier of the reader thread (thread-id or..). Henceforth, the first thread to taint a key is called the owner of that key. Once a key is tainted, it will never be written to by anyone other than the owner. Similarly, all consequent reads of a tainted key will also fail and only return the identifier of the owner (you could see where I might be going with this: transposition tables and task forwarding.) Only the owner will do further writes and reads, hence there is no more need for cross thread synchronisation.

What I ideally want here and I see missing is some mechanism that expresses: after the first tainting, no further operation should block. Both read and write should fail immediately. Hence, contention is sensible to exist only up until the first taint.

To the best of my knowledge, the known concurrent HashMap implementation (chashmap, evmap) will not work great here since they will not provide the exact functionality that I need. Aside from diving too deep into these, I first wanted to try and implement my own version of this, also to learn more about Rust's fearless concurrency.

I have two ideas to investigate more:

1. Using atomics

struct StateEntry<T> {
	data: T,
	taint: AtomicU8,
}

#[derive(Default, Debug)]
pub struct State {
	backend: RefCell<HashMap<Key, StateEntry<Value>>>,
}

And to taint, each thread will try and obtain the associated entry, and do an atomic compare and swap: if the previous entry.taint was 0 (the default, untainted), then swap it with my thread id. The rest of the operations will be done similarly and on top of this.

What I am uncertain here is if this is actually thread safe? sure, the CAS is atomic, but I am not sure what happens when two threads try and CAS into the same non-existent key. will it stay thread safe and atomic, or not? The code will look something like this:

fn try_taint(&self, key: Key, current: ThreadId) -> Result<(), ThreadId> {
	let owner = self
		.backend
		.borrow_mut()
		.entry(key)
		.or_default() 
		.taint // what if two threads execute up until here in parallel? will one of them be allowed to go forward?
		.compare_and_swap(0, current, SeqCst);
	if owner == current {
		Ok(())
	} else {
		Err(owner)
	}
}

2. Using RwLock only around the value, something along the lines of:

pub struct State {
	backend: RefCell<HashMap<Key, RwLock<StateEntry<Value>>>>,
}

Each thread will first try and acquire a write lock, then if failure fallback to a read lock (other way around could also work). I won't expand this further for brevity, but I think all in all it suffers from the same problem. What if two keys get initiated concurrently?

I think if an already existing key gets read concurrently it is fine, as there is some lock there and only one will be able to obtain it (or both of read request --anyhow), but the case of uninitialised entries is particularly unsettling for me.

---

Having written this, I think I am now certain that both approaches are wrong, since neither ensures that the same keys cannot be created at the same time. As far as I know, for the same reason, within the realm of the standard library, the common way to go is to wrap the whole HashMap inside a RwLock or Mutex. But this is surely too strict for my use case.

I guess my rephrased question would be what is the middle ground, if there is any? I precisely want to have locks around particular keys of the HashMap, not the whole thing at once, given the context provided for the tainting use case.

[raidwas]

Designing concurrent data structures is still hard, even if rust helps preventing errors.
Since you seem to only require a insert and lookup operation (no remove) there might be a better structure, but I think a concurrent hamt could be adapted to your needs.

[garagedragon]

Your first idea has a caveat that the thread that executes or_default is AFAIK not going to be guarenteed to be the one who get registers as owner in the CAS, but if that happens then the non-owner will still pick up on that correctly. However, the deeper problem with what you're doing is that AFAIK there's no safe way to get into that situation in the first place - RefCell is not Sync, so you won't be able to share it between threads, and the threadsafe version of the same functionality is... Mutex and RwLock

I precisely want to have locks around particular keys of the HashMap , not the whole thing at once , given the context provided for the tainting use case.

AFAIK the first idea of having an atomic "taint" against the value is workable. The standard Hashmap is Sync so long as both its keys and values (and internal state) are, so the problem for insertions is getting mutable access across threads. You can solve that by wrapping the map itself in a RwLock:

pub struct State {
	backend: RwLock<HashMap<Key, StateEntry<Value>>>,
}

This looks on its face like a non-starter because we don't want writes of existing values to contend with each other, but concurrent reads through a RwLock do not block each other, and you can take advantage of that by changing the map's value type:

struct StateEntry<T> {
	data: RefCell<T>,
	taint: AtomicU8,
}

This means we can make modifications through read locks in the RwLock. Unfortunately, it's difficult to construct and return references to the "guts" of the hashmap entries because they borrow local read handles, so the compiler objects. However, what you can do is take and call functions to operate on the values, which gives you this rather monsterous API:

type MapType = HashMap<Key, StateEntry<Value>>;
pub struct State {
    backend: RwLock<MapType>,
}
struct StateEntry<T> {
    data: RefCell<T>,
    taint: AtomicU8,
}


fn apply_mut<'a, T> (
    s: &State,
    k: Key,
    mut f: impl FnMut(&mut Value) -> T,
    thread_id: ThreadId,
) -> Result<Result<T, ThreadId>, TryLockError<RwLockReadGuard<'_, MapType>>> {
    match s.backend.try_read() {
        Ok(hm) => {
            let entry = hm.get(&k).unwrap();
            match entry.taint.load(Ordering::Relaxed) {
                taint if taint == thread_id => Ok(Ok(f(&mut (*entry.data.borrow_mut())))),
                taint => Ok(Err(taint))
            }
        }
        Err(e) => Err(e)
    }
}

The complexity of the return signature is mainly from how many things can go wrong in the process of trying to get hold of the StateEntry, but everything that can go wrong should do so immediately - the above function will not block unless the callback f does. (However it will panic if the entry doesn't exist)

AFAIK, to do inserts and removes you definitely will need to take write handles and lock out everyone else (who will return immediately with failure, not block) because you may be reallocating the map and there's no way to do that atomically. We can't do any tricks with or_default in the above because we don't hold a &mut HashMap at any point - the entry needs to exist for us to invoke the RefCell that gives us mutable references.

Hope that helps @kianenigma

[kianenigma]

@garagedragon, Thanks you so much for your info!

RefCell is not Sync , so you won't be able to share it between threads, and the threadsafe version of the same functionality is... Mutex and RwLock

Yeah this was my dumb mistake in the original snippet. The outer RefCell makes no sense at all. Indeed a RwLock should be the outer wrapper, at least it is the best that the standard library provides.

As for the rest of the description, what you present is exactly what I need: I don't want to acquire a write lock to write an already tainted value. And a RefCell, in principle, can solve that.

Unfortunately, it's difficult to construct and return references to the "guts" of the hashmap entries because they borrow local read handles, so the compiler objects

I haven't encountered the mutability issue that you mentioned, but have a rather more serious problem: If StateEntry contains RefCell (which as you said is not Sync), then all the types all the way up to State will not be Sync, hence they cannot be sticked into an Arc and shared across threads. I assume if you try and write a small test case for the above test case as well, you should get the same issue.

---

On a side note, I found this: atomic_refcell - Rust and the description is quite the much appealing and exactly what we're looking for here.

UPDATE: Indeed replacing RefCell with AtomicRefCell from the above crate makes the compiler happy. Yet, I have to read the code and comment of this crate to make sure its overhead is sensible, i.e. it does not add extra overhead when only one thread is writing to the value. Based on their initial explanation, this seems to be the case, so I am happy with it.

[garagedragon]

I believe, although not very confidentally, that you'd be safe to implement Sync directly on State and make the compiler happy that way. (Because we're manually making sure that only one thread ever sees any particular !Sync sub-field) An AtomicRefCell will cause a atomic mutation, which will AFAIK be very very slightly more expensive than an unsynchronised one, but if you're fine with that it's likely the safer solution