A new lock: AbLock

[xuxiaocheng0201]

I am building a complex resource lock that requires the use of this uncommon basic type of lock.

AbLock can be locked in two sides.
In each side, it can be locked many times. But cannot lock in both sides at the same time.

It sounds very simple, right?

But I can't implement it with regular Mutex and RwLock.
An obvious idea is two RwLock, read on the first and write on the second. But this way, it cannot be locked on the same side in many times.
So I think a lower implementation is needed, but I don't know how to do it specifically.

By the way, I need four methods: async lock_a_owned, async lock_b_owned, try_lock_a_owned and try_lock_b_owned. Basically corresponds to RwLock.

Can anyone provide implementation ideas? Thank you very much!

[Darksonn]

I think something along these lines would work:

use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;

pub struct AbLock {
    /// Bit 0 determines the kind of lock.
    /// Remaining bits count the number of times the lock has been taken.
    state: AtomicUsize,
    notify: [Notify; 2],
}

impl AbLock {
    pub fn try_lock(&self, kind: bool) -> Option<AbLockGuard<'_>> {
        let kind = kind as usize;
        let mut prev_state = self.state.load(Ordering::Relaxed);
        loop {
            let is_empty = prev_state <= 1;
            let is_right_kind = (prev_state & 1) == kind;

            let new_state = if is_empty {
                // Set counter to one. Set lower bit to the kind used to lock.
                2 | kind
            } else if is_right_kind {
                // increment counter by one
                prev_state.checked_add(2).unwrap()
            } else {
                // Lock is taken by the wrong kind.
                return None;
            };
            match self.state.compare_exchange_weak(
                prev_state,
                new_state,
                Ordering::AcqRel,
                Ordering::Relaxed,
            ) {
                Ok(_) => return Some(AbLockGuard(self)),
                Err(state) => prev_state = state,
            }
        }
    }

    pub async fn lock(&self, kind: bool) -> AbLockGuard<'_> {
        loop {
            let wait = self.notify[kind as usize].notified();

            if let Some(guard) = self.try_lock(kind) {
                return guard;
            }

            wait.await;
        }
    }

    fn unlock(&self) {
        let prev_state = self.state.fetch_sub(2, Ordering::AcqRel);
        let lock_kind = prev_state & 1;
        let num_locks = prev_state >> 1;
        if num_locks == 1 {
            // We just gave up the last lock.
            // Wake up waiters for the other kind of lock.
            self.notify[lock_kind ^ 1].notify_waiters();
        }
    }
}

pub struct AbLockGuard<'a>(&'a AbLock);

impl Drop for AbLockGuard<'_> {
    fn drop(&mut self) {
        self.0.unlock();
    }
}

[xuxiaocheng0201]

!!! Thank you very much!