mpscq.h

Go to the documentation of this file.

 
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#ifndef SBIT_MPSC_H_INCLUDED
#define SBIT_MPSC_H_INCLUDED
 
#include <sbit/stats.h>
 
#include <atomic>
#include <cstddef>
#include <list>
#include <memory_resource>
#include <vector>
 
namespace sbit::mpscq
{
 
class Queue
{
public:
   struct Envelope
   {
      Envelope* next;
 
      std::atomic<Envelope*>* workerPool;
 
      std::atomic<Envelope*>* ownerPool;
 
      void appendTo(std::atomic<Envelope*>* head)
      {
         next = head->load(std::memory_order_relaxed);
 
         while (!std::atomic_compare_exchange_weak_explicit(
            head, &next, this, std::memory_order_release, std::memory_order_relaxed))
         {
            // the body of the loop is empty
         }
      }
 
      void recycle()
      {
         appendTo(workerPool);
      }
   };
 
   template <typename Payload>
   struct Message
   {
      Envelope envelope;
 
      Payload payload;
 
      void recycle()
      {
         envelope.recycle();
      }
   };
 
   void append(Envelope* elem)
   {
      elem->appendTo(&m_head);
   }
 
   Envelope* flushAll()
   {
      Envelope* values = std::atomic_exchange_explicit(&m_head, nullptr, std::memory_order_release);
 
      // TODO(florin): reverse the links so we return a pointer to the oldest element
 
      return values;
   }
 
private:
   std::atomic<Envelope*> m_head{nullptr};
};
 
template <typename Payload>
class MessagePool
{
   using Envelope = Queue::Envelope;
   using Message  = Queue::Message<Payload>;
 
public:
   MessagePool(size_t allocationGroupSize, std::pmr::memory_resource* memoryResource) :
      m_allocationGroupSize{allocationGroupSize},
      m_recyclePool{nullptr},
      m_objectPool{allocationGroupSize, memoryResource}
   {
   }
 
   Message* allocate()
   {
      if (m_readyPool == nullptr)
      {
         m_readyPool = std::atomic_exchange_explicit(&m_recyclePool, nullptr, std::memory_order_release);
      }
 
      if (m_readyPool == nullptr)
      {
         try
         {
            auto& elems = m_objectPool.emplace_back(
               std::pmr::vector<Message>{m_allocationGroupSize, m_objectPool.get_allocator()});
            for (auto& msg : elems)
            {
               msg.envelope.next       = m_readyPool;
               m_readyPool             = &msg.envelope;
               msg.envelope.workerPool = &m_recyclePool;
               msg.envelope.ownerPool  = &m_recyclePool;
            }
         }
         catch (...)
         {
            return nullptr;
         }
      }
 
      auto* msg   = reinterpret_cast<Message*>(m_readyPool);
      m_readyPool = m_readyPool->next;
      return msg;
   }
 
private:
   const size_t m_allocationGroupSize;
 
   std::atomic<Envelope*> m_recyclePool;
 
   Envelope* m_readyPool = nullptr;
 
   std::pmr::list<std::pmr::vector<Message>> m_objectPool;
};
 
class ProcessorBase
{
public:
   explicit ProcessorBase(Queue& queue) : m_queue(queue)
   {
   }
 
   virtual void processElement(Queue::Envelope* envelope) = 0;
 
   virtual void afterBatch()
   {
      // do nothing
   }
 
   virtual void onIdle()
   {
      // do nothing
   }
 
   void startProcessing()
   {
      while (!m_done.load(std::memory_order_acquire))
      {
         auto* envelope = m_queue.flushAll();
 
         if (envelope == nullptr)
         {
            onIdle();
            continue;
         }
 
         while (envelope != nullptr)
         {
            processElement(envelope);
 
            auto* next = envelope->next;
            envelope->recycle();
            envelope = next;
         }
 
         afterBatch();
      }
   }
 
   void interrupt()
   {
      m_done.store(true, std::memory_order_release);
   }
 
   ProcessorBase(const ProcessorBase& other) = delete;
   ProcessorBase& operator=(const ProcessorBase& other) = delete;
 
   ProcessorBase(ProcessorBase&& other) = delete;
   ProcessorBase& operator=(ProcessorBase&& other) = delete;
 
   virtual ~ProcessorBase() = default;
 
private:
   Queue&            m_queue;
   std::atomic<bool> m_done{false};
};
 
template <typename Payload>
class Processor : public ProcessorBase
{
public:
   explicit Processor(Queue& queue) : ProcessorBase(queue)
   {
   }
 
protected:
   virtual void process(const Payload& payload) = 0;
 
private:
   using Message = Queue::Message<Payload>;
 
   void processElement(Queue::Envelope* envelope) override
   {
      auto* message = reinterpret_cast<Message*>(envelope);
      process(message->payload);
   }
};
 
} // namespace sbit::mpscq
 
#endif // SBIT_MPSC_H_INCLUDED