disrupt-rs 1.1.0

use crate::{
    barrier::{Barrier, NONE},
    consumer::Consumer,
    cursor::Cursor,
    producer::ProducerBarrier,
    producer::{Producer, RingBufferFull},
    ringbuffer::RingBuffer,
    wait_strategies::{WaitStrategy, WakeupNotifier},
    Sequence,
};
use crossbeam_utils::CachePadded;
use std::{
    hint, process,
    sync::{
        atomic::{fence, AtomicI64, AtomicU64, Ordering},
        Arc, Mutex,
    },
};

use super::{MissingFreeSlots, MutBatchIter};

type AvailableWord = CachePadded<AtomicU64>;

#[inline]
fn new_available_word(v: u64) -> AvailableWord {
    CachePadded::new(AtomicU64::new(v))
}

struct SharedProducer {
    consumers: Mutex<Vec<Consumer>>,
    ref_count: AtomicI64,
}

/// Producer for publishing to the Disruptor from multiple threads.
///
/// See also [SingleProducer](crate::SingleProducer) for single-threaded publication and
/// [`Producer`] for how to use a Producer.
pub struct MultiProducer<E, C, W>
where
    C: Barrier,
    W: WaitStrategy,
{
    shutdown_at_sequence: Arc<CachePadded<AtomicI64>>,
    ring_buffer: Arc<RingBuffer<E>>,
    shared_producer: Arc<SharedProducer>,
    producer_barrier: Arc<MultiProducerBarrier>,
    consumer_barrier: Arc<C>,
    notifier: W::Notifier,
    /// Next sequence number for this MultiProducer to publish.
    claimed_sequence: Sequence,
    /// Highest sequence available for publication because the Consumers are "enough" behind
    /// to not interfere.
    sequence_clear_of_consumers: Sequence,
}

impl<E, C, W> Producer<E> for MultiProducer<E, C, W>
where
    C: Barrier,
    W: WaitStrategy,
{
    #[inline]
    fn try_publish<F>(&mut self, update: F) -> Result<Sequence, RingBufferFull>
    where
        F: FnOnce(&mut E),
    {
        self.next_sequences(1).map_err(|_| RingBufferFull)?;
        let sequence = self.apply_update(update);
        Ok(sequence)
    }

    #[inline]
    fn publish<F>(&mut self, update: F)
    where
        F: FnOnce(&mut E),
    {
        while self.next_sequences(1).is_err() {
            hint::spin_loop();
        }
        self.apply_update(update);
    }

    #[inline]
    fn try_batch_publish<F>(&mut self, n: usize, update: F) -> Result<Sequence, MissingFreeSlots>
    where
        F: for<'a> FnOnce(MutBatchIter<'a, E>),
    {
        self.next_sequences(n)?;
        let sequence = self.apply_updates(n, update);
        Ok(sequence)
    }

    #[inline]
    fn batch_publish<F>(&mut self, n: usize, update: F)
    where
        F: for<'a> FnOnce(MutBatchIter<'a, E>),
    {
        debug_assert!(
            n <= self.ring_buffer.capacity() as usize,
            "batch_publish(n) requires n <= ring buffer capacity"
        );
        while self.next_sequences(n).is_err() {
            hint::spin_loop();
        }
        self.apply_updates(n, update);
    }
}

impl<E, C, W> Clone for MultiProducer<E, C, W>
where
    C: Barrier,
    W: WaitStrategy,
{
    fn clone(&self) -> Self {
        let count = self
            .shared_producer
            .ref_count
            .fetch_add(1, Ordering::AcqRel);

        // Cloning publishers and calling `mem::forget` on the clones could potentially overflow the
        // counter. It's very difficult to recover sensibly from such degenerate scenarios so we
        // just abort when the count becomes very large.
        if count > i64::MAX / 2 {
            process::abort();
        }

        let shutdown_at_sequence = Arc::clone(&self.shutdown_at_sequence);
        let producer_barrier = Arc::clone(&self.producer_barrier);
        let shared_producer = Arc::clone(&self.shared_producer);
        let consumer_barrier = Arc::clone(&self.consumer_barrier);
        let ring_buffer = Arc::clone(&self.ring_buffer);
        let notifier = self.notifier.clone();

        MultiProducer {
            shutdown_at_sequence,
            ring_buffer,
            shared_producer,
            producer_barrier,
            consumer_barrier,
            notifier,
            claimed_sequence: NONE,
            sequence_clear_of_consumers: 0,
        }
    }
}

impl<E, C, W> Drop for MultiProducer<E, C, W>
where
    C: Barrier,
    W: WaitStrategy,
{
    fn drop(&mut self) {
        let old_count = self
            .shared_producer
            .ref_count
            .fetch_sub(1, Ordering::AcqRel);
        if old_count == 1 {
            // All consumers are waiting to read the next sequence.
            let sequence = self.producer_barrier.current() + 1;
            self.shutdown_at_sequence.store(sequence, Ordering::Relaxed);
            // Wake any blocking consumers so they can observe shutdown promptly.
            self.notifier.wake();
            let mut consumers = self.shared_producer.consumers.lock().unwrap();
            consumers.iter_mut().for_each(|c| c.join());
        }
    }
}

impl<E, C, W> MultiProducer<E, C, W>
where
    C: Barrier,
    W: WaitStrategy,
{
    pub(crate) fn new(
        shutdown_at_sequence: Arc<CachePadded<AtomicI64>>,
        ring_buffer: Arc<RingBuffer<E>>,
        producer_barrier: Arc<MultiProducerBarrier>,
        consumers: Vec<Consumer>,
        consumer_barrier: C,
        notifier: W::Notifier,
    ) -> Self {
        let shared_producer = Arc::new(SharedProducer {
            consumers: Mutex::new(consumers),
            ref_count: AtomicI64::new(1),
        });
        let consumer_barrier = Arc::new(consumer_barrier);
        // Known to be available initially as consumers start at index 0.
        let sequence_clear_of_consumers = ring_buffer.capacity() - 1;
        MultiProducer {
            shutdown_at_sequence,
            ring_buffer,
            shared_producer,
            producer_barrier,
            consumer_barrier,
            notifier,
            claimed_sequence: NONE,
            sequence_clear_of_consumers,
        }
    }

    #[inline]
    fn next_sequences(&mut self, n: usize) -> Result<Sequence, MissingFreeSlots> {
        let n = n as i64;
        // We get the last produced sequence number and try and increment it.
        let mut current = self.producer_barrier.current();
        let mut n_next = current + n;

        loop {
            if self.sequence_clear_of_consumers < n_next {
                // We have to check where the rear consumer is in case we're about to
                // publish into the slot currently being read by the slowest consumer.
                // (Consumer is too far behind the producer to publish next n events).
                let rear_sequence_read = self.consumer_barrier.get_after(current);
                let available_capacity = self
                    .ring_buffer
                    .available_capacity(current, rear_sequence_read);
                if available_capacity < n {
                    return Err(MissingFreeSlots((n - available_capacity) as u64));
                }
                fence(Ordering::Acquire);

                // We now know how far we can continue until we get right behind the slowest consumers'
                // current position without checking where they actually are.
                self.sequence_clear_of_consumers = current + available_capacity;
            }

            match self.producer_barrier.compare_exchange(current, n_next) {
                Ok(_) => {
                    // The sequence interval `]current; n_next] is now exclusive for this producer.
                    self.claimed_sequence = n_next;
                    break;
                }
                Err(new_current) => {
                    current = new_current;
                    n_next = current + n;
                }
            }
        }

        Ok(n_next)
    }

    /// Precondition: `sequence` is available for publication.
    #[inline]
    fn apply_update<F>(&mut self, update: F) -> Sequence
    where
        F: FnOnce(&mut E),
    {
        let sequence = self.claimed_sequence;
        // SAFETY: Now, we have exclusive access to the event at `sequence` and a producer
        // can now update the data.
        let event_ptr = unsafe { self.ring_buffer.get(sequence).as_ptr() };
        let event = unsafe { &mut *event_ptr };
        update(event);
        // Make publication available by publishing `sequence`.
        self.producer_barrier.publish(sequence);
        self.notifier.wake();
        sequence
    }

    /// Precondition: `sequence` and previous `n - 1` sequences are available for publication.
    #[inline]
    fn apply_updates<F>(&mut self, n: usize, updates: F) -> Sequence
    where
        F: for<'a> FnOnce(MutBatchIter<'a, E>),
    {
        let n = n as i64;
        let upper = self.claimed_sequence;
        let lower = upper - n + 1;
        // SAFETY: Now, we have exclusive access to the event at `sequence` and a producer
        // can now update the data.
        let iter = MutBatchIter::new(lower, upper, &self.ring_buffer);
        updates(iter);
        // Make publications available by publishing all the sequences in the interval [lower; upper].
        fence(Ordering::Release);
        self.producer_barrier.publish_range_relaxed(lower, n);
        self.notifier.wake();
        self.claimed_sequence
    }
}

/// Barrier for multiple producers.
pub struct MultiProducerBarrier {
    /// Cursor used by producers to claim exclusive sequences per producer.
    cursor: Cursor,
    /// AtomicU64s each track availability of 64 slots.
    /// Each bit in the AtomicU64 encodes whether the slot was published in an even or odd round.
    /// This way producers avoid coordinating directly (with the added overhead that would have).
    ///
    /// Each word is cache-line padded to reduce false sharing between producers.
    /// Note, producers can never "overtake" each other and overwrite the availability of an event
    /// in the "previous" round as all producers must wait for the consumer furtherst behind (which
    /// is again blocked by the slowest producer).
    available: Box<[AvailableWord]>,
    index_mask: usize,
    index_shift: usize,
}

impl MultiProducerBarrier {
    pub(crate) fn new(size: usize) -> Self {
        assert!(
            size >= 64,
            "Multi Producer Disruptor must have a size of minimum 64 slots."
        );
        assert!(size.is_power_of_two(), "Size must be power of 2.");

        let cursor = Cursor::new(NONE);
        let i64_needed = size / 64;
        // available encodes 1 bit for each slot.
        // If the bit is 1 it means that the slot was published in the latest odd round
        // and 0 means latest even round. A ring buffer starts with round 0 (an even round)
        // so that is why we initialize with all 1's as nothing is published initially.
        let all_ones = !0_u64;
        let available = (0..i64_needed)
            .map(|_i| new_available_word(all_ones))
            .collect();
        let index_mask = size - 1;
        let index_shift = Self::log2(size);

        Self {
            cursor,
            available,
            index_mask,
            index_shift,
        }
    }

    fn log2(i: usize) -> usize {
        std::mem::size_of::<usize>() * 8 - (i.leading_zeros() as usize) - 1
    }

    #[inline]
    fn current(&self) -> Sequence {
        self.cursor.relaxed_value()
    }

    #[inline]
    fn compare_exchange(&self, current: Sequence, next: Sequence) -> Result<i64, i64> {
        self.cursor.compare_exchange(current, next)
    }

    /// Returns the availability index and the bit index of the given sequence number.
    #[inline]
    fn calculate_availability_indices(&self, sequence: Sequence) -> (usize, usize) {
        let slot_index = self.slot_index(sequence);
        let availability_index = slot_index >> 6; // == divide by 64.
        let bit_index = slot_index & 63;
        (availability_index, bit_index)
    }

    #[inline]
    fn slot_index(&self, sequence: Sequence) -> usize {
        sequence as usize & self.index_mask
    }

    /// Calculates if we're in an even (0) or odd (1) round.
    #[inline]
    fn calculate_availability_flag(&self, sequence: Sequence) -> u64 {
        let round = (sequence >> self.index_shift) as u64;
        round & 1
    }

    #[inline]
    fn availability_at(&self, index: usize) -> &AvailableWord {
        // SAFETY: Index is always calculated with `calculate_availability_index` and is therefore within bounds.
        unsafe { self.available.get_unchecked(index) }
    }

    #[inline]
    fn publish_range_relaxed(&self, from: Sequence, n: i64) {
        if n <= 0 {
            return;
        }

        let mut sequence = from;
        let mut remaining = n as usize;
        while remaining != 0 {
            let (availability_index, bit_index) = self.calculate_availability_indices(sequence);
            let take = remaining.min(64 - bit_index);

            let mask = if take == 64 {
                !0_u64
            } else {
                ((1_u64 << take) - 1) << bit_index
            };

            self.availability_at(availability_index)
                .fetch_xor(mask, Ordering::Relaxed);

            sequence += take as i64;
            remaining -= take;
        }
    }

    #[inline]
    fn publish_with_ordering(&self, sequence: Sequence, ordering: Ordering) {
        let (availability_index, bit_index) = self.calculate_availability_indices(sequence);
        let availability = self.availability_at(availability_index);
        let mask = 1 << bit_index;
        // XOR operation will flip the bit on exactly the bit_index position - encoding that we have
        // published an event in an even or odd round.
        availability.fetch_xor(mask, ordering);
    }
}

impl Barrier for MultiProducerBarrier {
    #[inline]
    fn get_after(&self, prev: Sequence) -> Sequence {
        let mut availability_flag = self.calculate_availability_flag(prev);
        let (mut availability_index, mut bit_index) = self.calculate_availability_indices(prev);
        let mut sequence = prev;

        loop {
            let word = self
                .availability_at(availability_index)
                .load(Ordering::Relaxed);
            let bits_remaining = 64 - bit_index;
            let relevant_mask = if bits_remaining == 64 {
                !0_u64
            } else {
                (1_u64 << bits_remaining) - 1
            };

            // Shift so bit 0 corresponds to `sequence`, then mask off the unused upper bits.
            let shifted = (word >> bit_index) & relevant_mask;
            let mismatch = if availability_flag == 0 {
                shifted
            } else {
                (!shifted) & relevant_mask
            };

            if mismatch != 0 {
                let matches = mismatch.trailing_zeros() as i64;
                return sequence + matches - 1;
            }

            // Entire remaining part of this word is contiguous and available.
            sequence += bits_remaining as i64;
            availability_index += 1;
            bit_index = 0;

            if availability_index == self.available.len() {
                availability_index = 0;
                availability_flag ^= 1;
            }
        }
    }
}

impl ProducerBarrier for MultiProducerBarrier {
    #[inline]
    fn publish(&self, sequence: Sequence) {
        self.publish_with_ordering(sequence, Ordering::Release);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn log2() {
        assert_eq!(1, MultiProducerBarrier::log2(2));
        assert_eq!(1, MultiProducerBarrier::log2(3));
        assert_eq!(3, MultiProducerBarrier::log2(8));
        assert_eq!(3, MultiProducerBarrier::log2(9));
        assert_eq!(3, MultiProducerBarrier::log2(10));
        assert_eq!(3, MultiProducerBarrier::log2(11));
    }

    #[test]
    fn initial_available_sequence_number() {
        let barrier = MultiProducerBarrier::new(64);

        assert_eq!(barrier.get_after(0), -1);
    }

    #[test]
    fn publication_of_single_event_for_small_barrier() {
        let barrier = MultiProducerBarrier::new(64);

        barrier.publish_range_relaxed(0, 1);
        // Verify published:
        assert_eq!(barrier.get_after(0), 0);
    }

    #[test]
    fn publication_of_range_for_small_barrier() {
        let barrier = MultiProducerBarrier::new(64);

        barrier.publish_range_relaxed(0, 10);
        // Verify published:
        assert_eq!(barrier.get_after(0), 9);
    }

    #[test]
    fn publication_of_range_wrapping_ringbuffer_for_small_barrier() {
        let barrier = MultiProducerBarrier::new(64);

        barrier.publish_range_relaxed(0, 50);
        // Verify published:
        assert_eq!(barrier.get_after(0), 49);

        barrier.publish_range_relaxed(50, 50);
        // Verify published:
        assert_eq!(barrier.get_after(49), 99);
    }

    #[test]
    fn publication_of_range_wrapping_ringbuffer_for_barrier() {
        let barrier = MultiProducerBarrier::new(128);

        barrier.publish_range_relaxed(0, 100);
        // Verify published:
        assert_eq!(barrier.get_after(0), 99);
        // Verify not published:

        barrier.publish_range_relaxed(100, 100);
        // Verify published:
        assert_eq!(barrier.get_after(99), 199);

        barrier.publish_range_relaxed(200, 100);
        // Verify published:
        assert_eq!(barrier.get_after(199), 299);
    }
}