Repository: qdrant/qdrant
File: lib/trififo/src/seqlock.rs (Rust)

Code before:

            let seq2 = self.seq.load(Ordering::Relaxed);

            // only return if seq did not change.
            if seq1 == seq2 {
                return result;
            }
        }
    }

    fn write(&self, callback: impl FnOnce(&mut T)) {
        let seq = self.seq.load(Ordering::Acquire);
        self.seq.store(seq + 1, Ordering::Release);

        fence(Ordering::Release);

        callback(unsafe { &mut *self.inner.get() });

        self.seq.store(seq + 2, Ordering::Relaxed);
    }
}

#[derive(Clone)]
pub struct SeqLockReader<T> {
    lock: Arc<SeqLock<T>>,
}

unsafe impl<T> Send for SeqLockReader<T> where T: Send {}
unsafe impl<T> Sync for SeqLockReader<T> {}

impl<T> SeqLockReader<T> {
    pub fn read<U>(&self, callback: impl Fn(&T) -> U) -> U {
        self.lock.read(callback)
    }
}

pub struct SeqLockWriter<T> {
    lock: Arc<SeqLock<T>>,
}

unsafe impl<T> Send for SeqLockWriter<T> where T: Send {}

impl<T> SeqLockWriter<T> {
    pub fn write(&self, callback: impl FnOnce(&mut T)) {
        self.lock.write(callback)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

Review comment (suggestion):
(about rust safety conventions)

This is extremely fragile.

Public methods `read` and `write` should be marked as `unsafe`. Or, alternatively, the callbacks should accept `*T`/`*mut T` instead of `&T`/`&mut T` to signify that these are not innocent references.

Either way, all callers should have `unsafe {}` block somewhere to be able to use these functions.

Code after:

    /// Wrap the value in a seqlock.
    ///
    /// Returns a reader which is Sync and Send; and a writer, which is only Send.
    ///
    /// The reader can be cloned and is safe to use between threads
    ///
    /// The writer should only have a single thread using it.
    pub fn new_reader_writer(v: T) -> (SeqLockReader<T>, SeqLockWriter<T>) {
        let lock = Arc::new(SeqLock::new(v));
        let reader = SeqLockReader { lock: lock.clone() };
        let writer = SeqLockWriter { lock: lock.clone() };
        (reader, writer)
    }

    fn read<U, F: Fn(&T) -> U>(&self, callback: F) -> U {
        loop {
            let seq1 = self.seq.load(Ordering::Acquire);

            // if odd, it is locked, wait until unlocked.
            if seq1 & 1 == 1 {
                std::hint::spin_loop();
                continue;
            }

            let result = callback(unsafe { &*self.inner.get() });

            // ensure that the read is complete before checking seq again
            fence(Ordering::Acquire);
            let seq2 = self.seq.load(Ordering::Relaxed);

            // only return if seq did not change.
            if seq1 == seq2 {
                return result;
            }
        }
    }
}

/// This structure can read the protected resource at any time, but it will:
/// - busy-wait if a write is taking place.
/// - or repeat the read if a write took place in between the start and end
///   of the read.
pub struct SeqLockReader<T: SeqLockSafe> {
    lock: Arc<SeqLock<T>>,
}

unsafe impl<T: Send + SeqLockSafe> Send for SeqLockReader<T> {}
unsafe impl<T: Sync + SeqLockSafe> Sync for SeqLockReader<T> {}

impl<T: SeqLockSafe> Clone for SeqLockReader<T> {