tor-browser

MultiWriterQueue.h (19910B)

---

/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#ifndef mozilla_MultiWriterQueue_h_
#define mozilla_MultiWriterQueue_h_

#include <cstdint>
#include <utility>

#include "RollingNumber.h"
#include "mozilla/Atomics.h"
#include "mozilla/MemoryReporting.h"
#include "mozilla/Mutex.h"
#include "prthread.h"

namespace mozilla {

// Default reader locking strategy, using a mutex to ensure that concurrent
// PopAll calls won't overlap.
class MOZ_CAPABILITY("mutex") MultiWriterQueueReaderLocking_Mutex {
public:
 MultiWriterQueueReaderLocking_Mutex()
     : mMutex("MultiWriterQueueReaderLocking_Mutex") {}
 void Lock() MOZ_CAPABILITY_ACQUIRE(mMutex) { mMutex.Lock(); };
 void Unlock() MOZ_CAPABILITY_RELEASE(mMutex) { mMutex.Unlock(); };

private:
 Mutex mMutex;
};

// Reader non-locking strategy, trusting that PopAll will never be called
// concurrently (e.g., by only calling it from a specific thread).
class MOZ_CAPABILITY("dummy lock") MultiWriterQueueReaderLocking_None {
public:
#ifndef DEBUG
 void Lock() MOZ_CAPABILITY_ACQUIRE() {};
 void Unlock() MOZ_CAPABILITY_RELEASE() {};
#else
 // DEBUG-mode checks to catch concurrent misuses.
 void Lock() MOZ_CAPABILITY_ACQUIRE() {
   MOZ_ASSERT(mLocked.compareExchange(false, true));
 };
 void Unlock() MOZ_CAPABILITY_RELEASE() {
   MOZ_ASSERT(mLocked.compareExchange(true, false));
 };

private:
 Atomic<bool> mLocked{false};
#endif
};

static constexpr uint32_t MultiWriterQueueDefaultBufferSize = 8192;

// Multi-writer, single-reader queue of elements of type `T`.
// Elements are bunched together in buffers of `BufferSize` elements.
//
// This queue is heavily optimized for pushing. In most cases pushes will only
// cost a couple of atomic reads and a few non-atomic reads. Worst cases:
// - Once per buffer, a push will allocate or reuse a buffer for later pushes;
// - During the above new-buffer push, other pushes will be blocked.
//
// By default, popping is protected by mutex; it may be disabled if popping is
// guaranteed never to be concurrent.
// In any case, popping will never negatively impact pushes.
// (However, *not* popping will add runtime costs, as unread buffers will not
// be freed, or made available to future pushes; Push functions provide
// feedback as to when popping would be most efficient.)
template <typename T, uint32_t BufferSize = MultiWriterQueueDefaultBufferSize,
         typename ReaderLocking = MultiWriterQueueReaderLocking_Mutex>
class MultiWriterQueue {
 static_assert(BufferSize > 0, "0-sized MultiWriterQueue buffer");

public:
 // Constructor.
 // Allocates the initial buffer that will receive the first `BufferSize`
 // elements. Also allocates one reusable buffer, which will definitely be
 // needed after the first `BufferSize` elements have been pushed.
 // Ideally (if the reader can process each buffer quickly enough), there
 // won't be a need for more buffer allocations.
 MultiWriterQueue()
     : mBuffersCoverAtLeastUpTo(BufferSize - 1),
       mMostRecentBuffer(new Buffer{}),
       mReusableBuffers(new Buffer{}),
       mOldestBuffer(static_cast<Buffer*>(mMostRecentBuffer)),
       mLiveBuffersStats(1),
       mReusableBuffersStats(1),
       mAllocatedBuffersStats(2) {}

 ~MultiWriterQueue() {
   auto DestroyList = [](Buffer* aBuffer) {
     while (aBuffer) {
       Buffer* older = aBuffer->Older();
       delete aBuffer;
       aBuffer = older;
     }
   };
   DestroyList(mMostRecentBuffer);
   DestroyList(mReusableBuffers);
 }

 // We need the index to be order-resistant to overflow, i.e., numbers before
 // an overflow should test smaller-than numbers after the overflow.
 // This is because we keep pushing elements with increasing Index, and this
 // Index is used to find the appropriate buffer based on a range; and this
 // need to work smoothly when crossing the overflow boundary.
 using Index = RollingNumber<uint32_t>;

 // Pushes indicate whether they have just reached the end of a buffer.
 using DidReachEndOfBuffer = bool;

 // Push new element and call aF on it.
 // Element may be in just-created state, or recycled after a PopAll call.
 // Atomically thread-safe; in the worst case some pushes may be blocked
 // while a new buffer is created/reused for them.
 // Returns whether that push reached the end of a buffer; useful if caller
 // wants to trigger processing regularly at the most efficient time.
 template <typename F>
 DidReachEndOfBuffer PushF(F&& aF) {
   // Atomically claim ownership of the next available element.
   const Index index{mNextElementToWrite++};
   // And now go and set that element.
   for (;;) {
     Index lastIndex{mBuffersCoverAtLeastUpTo};

     if (MOZ_UNLIKELY(index == lastIndex)) {
       // We have claimed the last element in the current head -> Allocate a
       // new head in advance of more pushes. Make it point at the current
       // most-recent buffer.
       // This whole process is effectively guarded:
       // - Later pushes will wait until mBuffersCoverAtLeastUpTo changes to
       //   one that can accept their claimed index.
       // - Readers will stop until the last element is marked as valid.
       Buffer* ourBuffer = mMostRecentBuffer;
       Buffer* newBuffer = NewBuffer(ourBuffer, index + 1);
       // Because we have claimed this very specific index, we should be the
       // only one touching the most-recent buffer pointer.
       MOZ_ASSERT(mMostRecentBuffer == ourBuffer);
       // Just pivot the most-recent buffer pointer to our new buffer.
       mMostRecentBuffer = newBuffer;
       // Because we have claimed this very specific index, we should be the
       // only one touching the buffer coverage watermark.
       MOZ_ASSERT(mBuffersCoverAtLeastUpTo == lastIndex.Value());
       // Update it to include the just-added most-recent buffer.
       mBuffersCoverAtLeastUpTo = index.Value() + BufferSize;
       // We know for sure that `ourBuffer` is the correct one for this index.
       ourBuffer->SetAndValidateElement(aF, index);
       // And indicate that we have reached the end of a buffer.
       return true;
     }

     if (MOZ_UNLIKELY(index > lastIndex)) {
       // We have claimed an element in a yet-unavailable buffer, wait for our
       // target buffer to be created (see above).
       while (Index(mBuffersCoverAtLeastUpTo) < index) {
         PR_Sleep(PR_INTERVAL_NO_WAIT);  // Yield
       }
       // Then loop to examine the new situation.
       continue;
     }

     // Here, we have claimed a number that is covered by current buffers.
     // These buffers cannot be destroyed, because our buffer is not filled
     // yet (we haven't written in it yet), therefore the reader thread will
     // have to stop there (or before) and won't destroy our buffer or more
     // recent ones.
     MOZ_ASSERT(index < lastIndex);
     Buffer* ourBuffer = mMostRecentBuffer;

     // In rare situations, another thread may have had the time to create a
     // new more-recent buffer, in which case we need to find our older buffer.
     while (MOZ_UNLIKELY(index < ourBuffer->Origin())) {
       // We assume that older buffers with still-invalid elements (e.g., the
       // one we have just claimed) cannot be destroyed.
       MOZ_ASSERT(ourBuffer->Older());
       ourBuffer = ourBuffer->Older();
     }

     // Now we can set&validate the claimed element, and indicate that we have
     // not reached the end of a buffer.
     ourBuffer->SetAndValidateElement(aF, index);
     return false;
   }
 }

 // Push new element and assign it a value.
 // Atomically thread-safe; in the worst case some pushes may be blocked
 // while a new buffer is created/reused for them.
 // Returns whether that push reached the end of a buffer; useful if caller
 // wants to trigger processing regularly at the most efficient time.
 DidReachEndOfBuffer Push(const T& aT) {
   return PushF([&aT](T& aElement, Index) { aElement = aT; });
 }

 // Push new element and move-assign it a value.
 // Atomically thread-safe; in the worst case some pushes may be blocked
 // while a new buffer is created/reused for them.
 // Returns whether that push reached the end of a buffer; useful if caller
 // wants to trigger processing regularly at the most efficient time.
 DidReachEndOfBuffer Push(T&& aT) {
   return PushF([&aT](T& aElement, Index) { aElement = std::move(aT); });
 }

 // Pop all elements before the first invalid one, running aF on each of them
 // in FIFO order.
 // Thread-safety with other PopAll calls is controlled by the `Locking`
 // template argument.
 // Concurrent pushes are always allowed, because:
 // - PopAll won't read elements until valid,
 // - Pushes do not interfere with pop-related members -- except for
 //   mReusableBuffers, which is accessed atomically.
 template <typename F>
 void PopAll(F&& aF) {
   mReaderLocking.Lock();
   // Destroy every second fully-read buffer.
   // TODO: Research a better algorithm, probably based on stats.
   bool destroy = false;
   for (;;) {
     Buffer* b = mOldestBuffer;
     MOZ_ASSERT(!b->Older());
     // The next element to pop must be in that oldest buffer.
     MOZ_ASSERT(mNextElementToPop >= b->Origin());
     MOZ_ASSERT(mNextElementToPop < b->Origin() + BufferSize);

     // Start reading each element.
     if (!b->ReadAndInvalidateAll(aF, mNextElementToPop)) {
       // Found an invalid element, stop popping.
       mReaderLocking.Unlock();
       return;
     }

     // Reached the end of this oldest buffer
     MOZ_ASSERT(mNextElementToPop == b->Origin() + BufferSize);
     // Delete this oldest buffer.
     // Since the last element was valid, it must mean that there is a newer
     // buffer.
     MOZ_ASSERT(b->Newer());
     MOZ_ASSERT(mNextElementToPop == b->Newer()->Origin());
     StopUsing(b, destroy);
     destroy = !destroy;

     // We will loop and start reading the now-oldest buffer.
   }
 }

 // Size of all buffers (used, or recyclable), excluding external data.
 size_t ShallowSizeOfExcludingThis(MallocSizeOf aMallocSizeOf) const {
   return mAllocatedBuffersStats.Count() * sizeof(Buffer);
 }

 struct CountAndWatermark {
   int mCount;
   int mWatermark;
 };

 CountAndWatermark LiveBuffersStats() const { return mLiveBuffersStats.Get(); }
 CountAndWatermark ReusableBuffersStats() const {
   return mReusableBuffersStats.Get();
 }
 CountAndWatermark AllocatedBuffersStats() const {
   return mAllocatedBuffersStats.Get();
 }

private:
 // Structure containing the element to be stored, and a validity-marker.
 class BufferedElement {
  public:
   // Run aF on an invalid element, and mark it as valid.
   template <typename F>
   void SetAndValidate(F&& aF, Index aIndex) {
     MOZ_ASSERT(!mValid);
     aF(mT, aIndex);
     mValid = true;
   }

   // Run aF on a valid element and mark it as invalid, return true.
   // Return false if element was invalid.
   template <typename F>
   bool ReadAndInvalidate(F&& aF) {
     if (!mValid) {
       return false;
     }
     aF(mT);
     mValid = false;
     return true;
   }

  private:
   T mT;
   // mValid should be atomically changed to true *after* mT has been written,
   // so that the reader can only see valid data.
   // ReleaseAcquire, because when set to `true`, we want the just-written mT
   // to be visible to the thread reading this `true`; and when set to `false`,
   // we want the previous reads to have completed.
   Atomic<bool, ReleaseAcquire> mValid{false};
 };

 // Buffer contains a sequence of BufferedElements starting at a specific
 // index, and it points to the next-older buffer (if any).
 class Buffer {
  public:
   // Constructor of the very first buffer.
   Buffer() : mOlder(nullptr), mNewer(nullptr), mOrigin(0) {}

   // Constructor of later buffers.
   Buffer(Buffer* aOlder, Index aOrigin)
       : mOlder(aOlder), mNewer(nullptr), mOrigin(aOrigin) {
     MOZ_ASSERT(aOlder);
     aOlder->mNewer = this;
   }

   Buffer* Older() const { return mOlder; }
   void SetOlder(Buffer* aOlder) { mOlder = aOlder; }

   Buffer* Newer() const { return mNewer; }
   void SetNewer(Buffer* aNewer) { mNewer = aNewer; }

   Index Origin() const { return mOrigin; }
   void SetOrigin(Index aOrigin) { mOrigin = aOrigin; }

   // Run aF on a yet-invalid element.
   // Not thread-safe by itself, but nothing else should write this element,
   // and reader won't access it until after it becomes valid.
   template <typename F>
   void SetAndValidateElement(F&& aF, Index aIndex) {
     MOZ_ASSERT(aIndex >= Origin());
     MOZ_ASSERT(aIndex < Origin() + BufferSize);
     mElements[aIndex - Origin()].SetAndValidate(aF, aIndex);
   }

   using DidReadLastElement = bool;

   // Read all valid elements starting at aIndex, marking them invalid and
   // updating aIndex.
   // Returns true if we ended up reading the last element in this buffer.
   // Accessing the validity bit is thread-safe (as it's atomic), but once
   // an element is valid, the reading itself is not thread-safe and should be
   // guarded.
   template <typename F>
   DidReadLastElement ReadAndInvalidateAll(F&& aF, Index& aIndex) {
     MOZ_ASSERT(aIndex >= Origin());
     MOZ_ASSERT(aIndex < Origin() + BufferSize);
     for (; aIndex < Origin() + BufferSize; ++aIndex) {
       if (!mElements[aIndex - Origin()].ReadAndInvalidate(aF)) {
         // Found an invalid element, stop here. (aIndex will not be updated
         // past it, so we will start from here next time.)
         return false;
       }
     }
     return true;
   }

  private:
   Buffer* mOlder;
   Buffer* mNewer;
   Index mOrigin;
   BufferedElement mElements[BufferSize];
 };

 // Reuse a buffer, or create a new one.
 // All buffered elements will be invalid.
 Buffer* NewBuffer(Buffer* aOlder, Index aOrigin) {
   MOZ_ASSERT(aOlder);
   for (;;) {
     Buffer* head = mReusableBuffers;
     if (!head) {
       ++mAllocatedBuffersStats;
       ++mLiveBuffersStats;
       Buffer* buffer = new Buffer(aOlder, aOrigin);
       return buffer;
     }
     Buffer* older = head->Older();
     // Try to pivot the reusable-buffer pointer from the current head to the
     // next buffer in line.
     if (mReusableBuffers.compareExchange(head, older)) {
       // Success! The reusable-buffer pointer now points at the older buffer,
       // so we can recycle this ex-head.
       --mReusableBuffersStats;
       ++mLiveBuffersStats;
       head->SetOlder(aOlder);
       aOlder->SetNewer(head);
       // We will be the newest; newer-pointer should already be null.
       MOZ_ASSERT(!head->Newer());
       head->SetOrigin(aOrigin);
       return head;
     }
     // Failure, someone else must have touched the list, loop to try again.
   }
 }

 // Discard a fully-read buffer.
 // If aDestroy is true, delete it.
 // If aDestroy is false, move the buffer to a reusable-buffer stack.
 void StopUsing(Buffer* aBuffer, bool aDestroy) {
   --mLiveBuffersStats;

   // We should only stop using the oldest buffer.
   MOZ_ASSERT(!aBuffer->Older());
   // The newest buffer should not be modified here.
   MOZ_ASSERT(aBuffer->Newer());
   MOZ_ASSERT(aBuffer->Newer()->Older() == aBuffer);
   // Detach from the second-oldest buffer.
   aBuffer->Newer()->SetOlder(nullptr);
   // Make the second-oldest buffer the now-oldest buffer.
   mOldestBuffer = aBuffer->Newer();

   if (aDestroy) {
     --mAllocatedBuffersStats;
     delete aBuffer;
   } else {
     ++mReusableBuffersStats;
     // The recycling stack only uses mOlder; mNewer is not needed.
     aBuffer->SetNewer(nullptr);

     // Make the given buffer the new head of reusable buffers.
     for (;;) {
       Buffer* head = mReusableBuffers;
       aBuffer->SetOlder(head);
       if (mReusableBuffers.compareExchange(head, aBuffer)) {
         break;
       }
     }
   }
 }

 // Index of the next element to write. Modified when an element index is
 // claimed for a push. If the last element of a buffer is claimed, that push
 // will be responsible for adding a new head buffer.
 // Relaxed, because there is no synchronization based on this variable, each
 // thread just needs to get a different value, and will then write different
 // things (which themselves have some atomic validation before they may be
 // read elsewhere, independent of this `mNextElementToWrite`.)
 Atomic<Index::ValueType, Relaxed> mNextElementToWrite{0};

 // Index that a live recent buffer reaches. If a push claims a lesser-or-
 // equal number, the corresponding buffer is guaranteed to still be alive:
 // - It will have been created before this index was updated,
 // - It will not be destroyed until all its values have been written,
 //   including the one that just claimed a position within it.
 // Also, the push that claims this exact number is responsible for adding the
 // next buffer and updating this value accordingly.
 // ReleaseAcquire, because when set to a certain value, the just-created
 // buffer covering the new range must be visible to readers.
 Atomic<Index::ValueType, ReleaseAcquire> mBuffersCoverAtLeastUpTo;

 // Pointer to the most recent buffer. Never null.
 // This is the most recent of a deque of yet-unread buffers.
 // Only modified when adding a new head buffer.
 // ReleaseAcquire, because when modified, the just-created new buffer must be
 // visible to readers.
 Atomic<Buffer*, ReleaseAcquire> mMostRecentBuffer;

 // Stack of reusable buffers.
 // ReleaseAcquire, because when modified, the just-added buffer must be
 // visible to readers.
 Atomic<Buffer*, ReleaseAcquire> mReusableBuffers;

 // Template-provided locking mechanism to protect PopAll()-only member
 // variables below.
 ReaderLocking mReaderLocking;

 // Pointer to the oldest buffer, which contains the new element to be popped.
 // Never null.
 Buffer* mOldestBuffer;

 // Index of the next element to be popped.
 Index mNextElementToPop{0};

 // Stats.
 class AtomicCountAndWatermark {
  public:
   explicit AtomicCountAndWatermark(int aCount)
       : mCount(aCount), mWatermark(aCount) {}

   int Count() const { return int(mCount); }

   CountAndWatermark Get() const {
     return CountAndWatermark{int(mCount), int(mWatermark)};
   }

   int operator++() {
     int count = int(++mCount);
     // Update watermark.
     for (;;) {
       int watermark = int(mWatermark);
       if (watermark >= count) {
         // printf("++[%p] -=> %d-%d\n", this, count, watermark);
         break;
       }
       if (mWatermark.compareExchange(watermark, count)) {
         // printf("++[%p] -x> %d-(was %d now %d)\n", this, count, watermark,
         // count);
         break;
       }
     }
     return count;
   }

   int operator--() {
     int count = int(--mCount);
     // printf("--[%p] -> %d\n", this, count);
     return count;
   }

  private:
   // Relaxed, as these are just gathering stats, so consistency is not
   // critical.
   Atomic<int, Relaxed> mCount;
   Atomic<int, Relaxed> mWatermark;
 };
 // All buffers in the mMostRecentBuffer deque.
 AtomicCountAndWatermark mLiveBuffersStats;
 // All buffers in the mReusableBuffers stack.
 AtomicCountAndWatermark mReusableBuffersStats;
 // All allocated buffers (sum of above).
 AtomicCountAndWatermark mAllocatedBuffersStats;
};

}  // namespace mozilla

#endif  // mozilla_MultiWriterQueue_h_