ringbuffer.py

"""Ring buffers for multiprocessing.

Allows multiple child Python processes started via the multiprocessing module
to read from a shared ring buffer in the parent process. For each child, a
pointer is maintained for the purpose of reading. One pointer is maintained by
for the purpose of writing. Reads may be issued in blocking or non-blocking
mode. Writes are always in non-blocking mode and will raise an exception
if the buffer is full.
"""

import ctypes
import contextlib
import functools
import multiprocessing
import struct


class Error(Exception):
    pass


class DataTooLargeError(Error, ValueError):
    pass


class WaitingForReaderError(Error):
    pass


class WaitingForWriterError(Error):
    pass


class WriterFinishedError(Error):
    pass


class AlreadyClosedError(Error):
    pass


class MustCreatedReadersBeforeWritingError(Error):
    pass


class InternalLockingError(Error):
    pass


class Position:

    def __init__(self, slot_count):
        self.counter = 0
        self.slot_count = slot_count

    @property
    def index(self):
        return self.counter % self.slot_count

    @property
    def generation(self):
        return self.counter // self.slot_count


class Pointer:

    def __init__(self, slot_count, *, start=None):
        default = start if start is not None else 0
        self.counter = multiprocessing.RawValue(ctypes.c_longlong, default)
        self.position = Position(slot_count)

    def increment(self):
        self.counter.value += 1

    def get(self):
        # Avoid reallocating Position repeatedly.
        self.position.counter = self.counter.value
        return self.position

    def set(self, counter):
        self.counter.value = counter


class RingBuffer:
    """Circular buffer class accessible to multiple threads or child processes.

    All methods are thread safe. Multiple readers and writers are permitted.
    Before kicking off multiprocessing.Process instances, first allocate all
    of the writers you'll need with new_writer() and readers with new_reader().
    Pass the Pointer value returned by the new_reader() method to the
    multiprocessing.Process constructor along with the RingBuffer instance.
    Calling  new_writer() or new_reader() from a child multiprocessing.Process
    will not work.
    """

    def __init__(self, *, slot_bytes, slot_count):
        """Initializer.

        Args:
            slot_bytes: The maximum size of slots in the buffer.
            slot_count: How many slots should be in the buffer.
        """
        self.slot_count = slot_count
        self.array = SlotArray(slot_bytes=slot_bytes, slot_count=slot_count)
        self.lock = ReadersWriterLock()
        # Each reading process may modify its own Pointer while the read
        # lock is being held. Each reading process can also load the position
        # of the writer, but not load any other readers. Each reading process
        # can also load the value of the 'active' count.
        self.readers = []
        # The writer can load and store the Pointer of all the reader Pointers
        # or the writer Pointer while the write lock is held. It can also load
        # and store the value of the 'active' acount.
        self.writer = Pointer(self.slot_count)
        self.active = multiprocessing.RawValue(ctypes.c_uint, 0)

    def new_reader(self):
        """Returns a new unique reader into the buffer.

        This must only be called in the parent process. It must not be
        called in a child multiprocessing.Process. See class docstring. To
        enforce this policy, no readers may be allocated after the first
        write has occurred.
        """
        with self.lock.for_write():
            writer_position = self.writer.get()
            if writer_position.counter > 0:
                raise MustCreatedReadersBeforeWritingError

            reader = Pointer(self.slot_count, start=writer_position.counter)
            self.readers.append(reader)
            return reader

    def new_writer(self):
        """Must be called once by each writer before any reads occur.

        Should be paired with a single subsequent call to writer_done() to
        indicate that this writer has finished and will not write any more
        data into the ring.
        """
        with self.lock.for_write():
            self.active.value += 1

    def _has_write_conflict(self, position):
        index = position.index
        generation = position.generation
        for reader in self.readers:
            # This Position and the other Position both point at the same index
            # in the ring buffer, but they have different generation numbers.
            # This means the writer can't proceed until some readers have
            # sufficiently caught up.
            reader_position = reader.get()
            if (reader_position.index == index and
                    reader_position.generation < generation):
                return True

        return False

    def try_write(self, data):
        """Tries to write the next slot, but will not block.

        Once a successful write occurs, all pending blocking_read() calls
        will be woken up to consume the newly written slot.

        Args:
            data: Bytes to write in the next available slot. Must be
                less than or equal to slot_bytes in size.

        Raises:
            WaitingForReaderError: If all of the slots are full and we need
                to wait for readers to catch up before there will be
                sufficient room to write more data. This is a sign that
                the readers can't keep up with the writer. Consider calling
                force_reader_sync() if you need to force the readers to
                catch up, but beware that means they will miss data.
        """
        with self.lock.for_write():
            if not self.active.value:
                raise AlreadyClosedError

            position = self.writer.get()
            if self._has_write_conflict(position):
                raise WaitingForReaderError

            self.array[position.index] = data
            self.writer.increment()

    def _has_read_conflict(self, reader_position):
        writer_position = self.writer.get()
        return writer_position.counter <= reader_position.counter

    def _try_read_no_lock(self, reader):
        position = reader.get()
        if self._has_read_conflict(position):
            if not self.active.value:
                raise WriterFinishedError
            else:
                raise WaitingForWriterError

        data = self.array[position.index]
        reader.increment()
        return data

    def try_read(self, reader):
        """Tries to read the next slot for a reader, but will not block.

        Args:
            reader: Position previously returned by the call to new_reader().

        Returns:
            bytearray containing a copy of the data from the slot. This
            value is mutable an can be used to back ctypes objects, NumPy
            arrays, etc.

        Raises:
            WriterFinishedError: If the RingBuffer was closed before this
                read operation began.
            WaitingForWriterError: If the given reader has already consumed
                all the data in the ring buffer and would need to block in
                order to wait for new data to arrive.
        """
        with self.lock.for_read():
            return self._try_read_no_lock(reader)

    def blocking_read(self, reader):
        """Reads the next slot for a reader, blocking if it isn't filled yet.

        Args:
            reader: Position previously returned by the call to new_reader().

        Returns:
            bytearray containing a copy of the data from the slot. This
            value is mutable an can be used to back ctypes objects, NumPy
            arrays, etc.

        Raises:
            WriterFinishedError: If the RingBuffer was closed while waiting
                to read the next operation.
        """
        with self.lock.for_read():
            while True:
                try:
                    return self._try_read_no_lock(reader)
                except WaitingForWriterError:
                    self.lock.wait_for_write()

    def force_reader_sync(self):
        """Forces all readers to skip to the position of the writer."""
        with self.lock.for_write():
            writer_position = self.writer.get()

            for reader in self.readers:
                reader.set(writer_position.counter)

            for reader in self.readers:
                p = reader.get()

    def force_single_reader_sync(self, reader):
        """Forces passed reader to skip to the position of the writer."""
        with self.lock.for_write():
            writer_position = self.writer.get()

            reader.set(writer_position.counter)

            p = reader.get()

    def writer_done(self):
        """Called by the writer when no more data is expected to be written.

        Should be called once for every corresponding call to new_writer().
        Once all writers have called writer_done(), a WriterFinishedError
        exception will be raised by any blocking read calls or subsequent
        calls to read.
        """
        with self.lock.for_write():
            self.active.value -= 1


class SlotArray:
    """Fast array of indexable buffers backed by shared memory.

    Assumes locking happens elsewhere.
    """

    def __init__(self, *, slot_bytes, slot_count):
        """Initializer.

        Args:
            slot_bytes: How big each buffer in the array should be.
            slot_count: How many buffers should be in the array.
        """
        self.slot_bytes = slot_bytes
        self.slot_count = slot_count
        self.length_bytes = 4
        slot_type = ctypes.c_byte * (slot_bytes + self.length_bytes)
        self.array = multiprocessing.RawArray(slot_type, slot_count)

    def __getitem__(self, i):
        data = memoryview(self.array[i])
        (length,) = struct.unpack_from('>I', data, 0)

        start = self.length_bytes
        # This must create a copy because we want the writer to be able to
        # overwrite this slot as soon as the data has been retrieved by all
        # readers. But we also want the returned bytes to be mutable so that
        # the returned data can immediately back a ctypes record using the
        # from_buffer() method (instead of from_buffer_copy()).
        return bytearray(data[start:start + length])

    def __setitem__(self, i, data):
        data_view = memoryview(data).cast('@B')
        data_size = len(data_view)
        if data_size > self.slot_bytes:
            raise DataTooLargeError('%d bytes too big for slot' % data_size)

        # Avoid copying the input data! Do only a single copy into the slot.
        slot_view = memoryview(self.array[i]).cast('@B')
        struct.pack_into('>I', slot_view, 0, data_size)
        start = self.length_bytes
        slot_view[start:start + data_size] = data_view

    def __len__(self):
        return self.slot_count


class ReadersWriterLock:
    """Multiprocessing-compatible Readers/Writer lock.

    The algorithm:
    https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Using_a_condition_variable_and_a_mutex

    Background on the Kernel:
    https://www.kernel.org/doc/Documentation/memory-barriers.txt

    sem_wait on Linux uses NPTL, which uses futexes:
    https://github.com/torvalds/linux/blob/master/kernel/futex.c

    Notably, futexes use the smp_mb() memory fence, which is a general write
    barrier, meaning we can assume that all memory reads and writes before
    a barrier will complete before reads and writes after the barrier, even
    if the semaphore / futex isn't actively held.
    """

    def __init__(self):
        self.lock = multiprocessing.Lock()
        self.readers_condition = multiprocessing.Condition(self.lock)
        self.writer_condition = multiprocessing.Condition(self.lock)
        self.readers = multiprocessing.RawValue(ctypes.c_uint, 0)
        self.writer = multiprocessing.RawValue(ctypes.c_bool, False)

    def _acquire_reader_lock(self):
        with self.lock:
            while self.writer.value:
                self.readers_condition.wait()

            self.readers.value += 1

    def _release_reader_lock(self):
        with self.lock:
            self.readers.value -= 1

            if self.readers.value == 0:
                self.writer_condition.notify()

    @contextlib.contextmanager
    def for_read(self):
        """Acquire the lock for reading."""
        self._acquire_reader_lock()
        try:
            yield
        finally:
            self._release_reader_lock()

    def _acquire_writer_lock(self):
        with self.lock:
            while self.writer.value or self.readers.value > 0:
                self.writer_condition.wait()

            self.writer.value = True

    def _release_writer_lock(self):
        with self.lock:
            self.writer.value = False
            self.readers_condition.notify_all()
            self.writer_condition.notify()

    @contextlib.contextmanager
    def for_write(self):
        """Acquire the lock for writing reading."""
        self._acquire_writer_lock()
        try:
            yield
        finally:
            self._release_writer_lock()

    def wait_for_write(self):
        """Block until a writer has notified readers.

        Must be called while the read lock is already held. May return
        spuriously before the writer actually did something.
        """
        with self.lock:
            if self.readers.value == 0:
                raise InternalLockingError
            self.readers.value -= 1
            self.writer_condition.notify()
            self.readers_condition.wait()
            self.readers.value += 1