pub_sub.hpp

Go to the documentation of this file.

/*
 
Copyright 2025 Matthew Tolman
 
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 MTCORE_THREAD_PUB_SUB_HPP
#define MTCORE_THREAD_PUB_SUB_HPP
 
#include "mtcore/colls/gen_list.hpp"
#include "mtcore_thread/arc.hpp"
#include "mtcore_thread/core.hpp"
#include "mtcore/alloc.hpp"
#include "mtcore_thread/traits.hpp"
#include <memory>
 
namespace mtcore::thread {

  template<typename T>
  struct Subscription;
 
  namespace impl {
    template<typename T>
    struct Broker {
      Allocator *alloc;
      size_t queueSize;
      std::timed_mutex mux;
      std::condition_variable_any readSignal;
      std::condition_variable_any senderSignal;
      bool closed            = false;
      bool hasReservedWriter = false;
 
      Broker(Allocator *alloc, size_t queueSize) : alloc(alloc), queueSize(queueSize) { ensure(alloc); }
 
      GenList<Queue<T>> inboxes;
 
      ~Broker() {
        auto lock = std::unique_lock(mux);
        auto iter = inboxes.ptr_iter();
        Queue<T> *q;
        while (iter.next().copy_if_present(q)) {
          print_debug("%s", "NOT ALL BROKER SUBSCRIPTIONS WERE CLEANED UP! CLEANING UP A "
                            "SUBSCRIPTION IN Arc DESTRUCTOR!");
          q->deinit(*alloc);
        }
 
        inboxes.deinit(*alloc);
      }
    };
  }  // namespace impl

  enum class PubSubSendBlockGrowErrors { CHANNEL_CLOSED, PARTIAL_SEND, ALLOCATION_FAILED };
  static_assert(has_closed_error<PubSubSendBlockGrowErrors>);

  enum class PubSubSendBeforeGrowsError { CHANNEL_CLOSED, TIMEOUT, PARTIAL_SEND, ALLOCATION_FAILED };
  static_assert(has_closed_error<PubSubSendBeforeGrowsError>);

  enum class PubSubTrySendGrowsError { CHANNEL_CLOSED, COULD_NOT_LOCK, NO_ROOM, PARTIAL_SEND, ALLOCATION_FAILED };
  static_assert(has_closed_error<PubSubTrySendGrowsError>);
 
  template<typename T>
  struct Subscription {
    using Message = T;

    Result<Optional<T>, ChannelTryReceiveErrors> try_receive() {
      if (!broker.valid()) {
        return success(Optional<T>{nullopt});
      }
 
      auto lock = std::unique_lock(broker->mux, std::try_to_lock_t{});
      if (!lock.owns_lock()) {
        return error(ChannelTryReceiveErrors::COULD_NOT_LOCK);
      }
 
      if (!broker->inboxes.is_handle_valid(inboxHandle)) {
        return success(Optional<T>{nullopt});
      }
 
      if (auto res = broker->inboxes[inboxHandle].pop(); res.has_value()) {
        broker->senderSignal.notify_all();
        return res;
      }
 
      if (broker->closed) {
        return success(Optional<T>{nullopt});
      }
 
      return error(ChannelTryReceiveErrors::MESSAGE_NOT_READY);
    }

    Optional<T> receive_block() {
      if (!broker.valid()) {
        return Optional<T>{nullopt};
      }
 
      auto lock = std::unique_lock(broker->mux);
 
      if (!broker->inboxes.is_handle_valid(inboxHandle)) {
        return Optional<T>{nullopt};
      }
 
      do {
        if (auto res = broker->inboxes[inboxHandle].pop(); res.has_value()) {
          broker->senderSignal.notify_all();
          return res;
        }
 
        if (broker->closed) {
          return Optional<T>{nullopt};
        }
        broker->readSignal.wait(lock);
      } while (true);
    }

    template<typename Clock, typename Duration>
    Result<Optional<T>, ChannelReceiveBeforeErrors>
    receive_before(const std::chrono::time_point<Clock, Duration> &timePoint) {
      if (!broker.valid()) {
        return success(Optional<T>{nullopt});
      }
 
      auto lock = std::unique_lock(broker->mux, timePoint);
      if (!lock.owns_lock()) {
        return error(ChannelReceiveBeforeErrors::TIMEOUT);
      }
 
      if (!broker->inboxes.is_handle_valid(inboxHandle)) {
        return success(Optional<T>{nullopt});
      }
 
      do {
        if (auto res = broker->inboxes[inboxHandle].pop(); res.has_value()) {
          broker->senderSignal.notify_all();
          return res;
        }
 
        if (broker->closed) {
          return success(Optional<T>{nullopt});
        }
 
        if (broker->readSignal.wait_until(lock, timePoint) == std::cv_status::timeout) {
          return error(ChannelReceiveBeforeErrors::TIMEOUT);
        }
      } while (true);
    }

    void deinit(Allocator &alloc) {
      if (!broker.valid()) {
        return;
      }
      {
        auto lock = std::unique_lock(broker->mux);
        broker->inboxes[inboxHandle].deinit(*broker->alloc);
        broker->inboxes.remove(inboxHandle);
        broker->senderSignal.notify_all();
      }
      broker.release(alloc);
    }
 
    void close() { deinit(); }
 
    Subscription() = default;
 
    Subscription(const Arc<impl::Broker<T>> &broker, const Handle &handle) : broker(broker), inboxHandle(handle) {}
 
  private:
    Arc<impl::Broker<T>> broker;
    Handle inboxHandle = {0, 0};
  };
 
  static_assert(is_receiver<Subscription<int>>, "SUBSCRIPTION ISN'T A RECEIVER");
  static_assert(is_receiver_block<Subscription<int>>, "SUBSCRIPTION ISN'T A BLOCK RECEIVER");
  static_assert(is_receiver_before<Subscription<int>>, "SUBSCRIPTION ISN'T A BEFORE RECEIVER");
  static_assert(is_closeable<Subscription<int>>, "SUBSCRIPTION ISN'T A CLOSEABLE");

  template<typename T>
  struct Publisher {
    using Message = T;
 
    Arc<impl::Broker<T>> broker;

    Result<void, AllocationError> init(Allocator &alloc, size_t queueSize = 10) {
      ensure(!broker.valid(), "ALREADY INITIALIZED");
      if (auto brokerRes = broker.init(alloc, &alloc, queueSize); brokerRes.is_error()) {
        return brokerRes.error();
      }
      if (auto res = broker->inboxes.init(alloc); res.is_error()) {
        broker.deinit(alloc);
        return res.error();
      }
      return success();
    }

    void close() {
      auto lock = std::unique_lock(broker->mux);
 
      // close out all our inboxes
      broker->closed = true;
      broker->senderSignal.notify_all();
      broker->readSignal.notify_all();
    }

    void deinit(Allocator &alloc) {
      if (!broker.valid()) {
        return;
      }
      close();
      broker.deinit(alloc);
    }

    Result<Subscription<T>, AllocationError> subscribe() {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux);
      Queue<T> q;
      if (auto res = q.init(*broker->alloc, broker->queueSize); res.is_error()) {
        return res.error();
      }
 
      auto handleRes = broker->inboxes.add(q);
      if (handleRes.is_error()) {
        q.deinit(*broker->alloc);
        return error(AllocationError::ALLOCATION_FAILED);
      }
 
      auto subRef = broker.acquire(*broker->alloc);
      if (subRef.is_error()) {
        broker->inboxes.remove(handleRes.value());
        q.deinit(*broker->alloc);
        return error(AllocationError::ALLOCATION_FAILED);
      }
      return success(Subscription<T>{subRef.value(), *handleRes});
    }

    Result<void, ChannelSendBlockErrors> send_block(const T &message) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux);
 
      if (broker->closed) {
        return error(ChannelSendBlockErrors::CHANNEL_CLOSED);
      }
 
      bool ready;
      do {
        ready = true;
        // pre-check that we won't error out when adding
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          if (cur->is_full()) {
            ready = false;
          }
        }
        if (!ready) {
          broker->senderSignal.wait(lock);
        }
      } while (!ready);
 
      {
        // Now that we're good, send
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          ensure(cur->push(message).is_success(), "PRE CHECK FAILED");
        }
        broker->readSignal.notify_all();
      }
      return success();
    }

    Result<void, PubSubSendBlockGrowErrors> send_block(Allocator &alloc, const T &message) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux);
 
      if (broker->closed) {
        return error(PubSubSendBlockGrowErrors::CHANNEL_CLOSED);
      }
 
      auto iter = broker->inboxes.ptr_iter();
      Queue<T> *cur;
      bool allocated                              = false;
      bool sentSomething                          = false;
      Result<void, PubSubSendBlockGrowErrors> res = success();
      while (iter.next().copy_if_present(cur)) {
        allocated |= cur->is_full();
        if (!cur->push(alloc, message)) {
          res = error(PubSubSendBlockGrowErrors::ALLOCATION_FAILED);
        }
        else {
          sentSomething = true;
        }
      }
 
      if (res.is_error() && sentSomething) {
        res = error(PubSubSendBlockGrowErrors::PARTIAL_SEND);
      }
 
      if (sentSomething) {
        broker->readSignal.notify_all();
        if (allocated) {
          broker->senderSignal.notify_all();
        }
      }
      return res;
    }

    template<typename Clock, typename Duration>
    Result<void, ChannelSendBeforeErrors> send_before(const T &message,
                                                      const std::chrono::time_point<Clock, Duration> &timePoint) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux, timePoint);
      if (!lock.owns_lock()) {
        return error(ChannelSendBeforeErrors::TIMEOUT);
      }
 
      if (broker->closed) {
        return error(ChannelSendBeforeErrors::CHANNEL_CLOSED);
      }
 
      bool ready;
      do {
        ready = true;
        // pre-check that we won't error out when adding
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          if (cur->is_full()) {
            ready = false;
          }
        }
        if (!ready) {
          if (broker->senderSignal.wait_until(lock, timePoint) == std::cv_status::timeout) {
            return error(ChannelSendBeforeErrors::TIMEOUT);
          }
        }
      } while (!ready);
      {
        // Now that we're good, send
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          ensure(cur->push(message).is_success(), "PRE CHECK FAILED");
        }
        broker->readSignal.notify_all();
      }
      return success();
    }

    template<typename Clock, typename Duration>
    Result<void, PubSubSendBeforeGrowsError> send_before(Allocator &alloc, const T &message,
                                                         const std::chrono::time_point<Clock, Duration> &timePoint) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux, timePoint);
      if (!lock.owns_lock()) {
        return error(PubSubSendBeforeGrowsError::TIMEOUT);
      }
 
      if (broker->closed) {
        return error(PubSubSendBeforeGrowsError::CHANNEL_CLOSED);
      }
 
      auto iter = broker->inboxes.ptr_iter();
      Queue<T> *cur;
      bool allocated                               = false;
      bool sentSomething                           = false;
      Result<void, PubSubSendBeforeGrowsError> res = success();
      while (iter.next().copy_if_present(cur)) {
        allocated |= cur->is_full();
        if (!cur->push(alloc, message)) {
          res = error(PubSubSendBeforeGrowsError::ALLOCATION_FAILED);
        }
        else {
          sentSomething = true;
        }
      }
 
      if (res.is_error() && sentSomething) {
        res = error(PubSubSendBeforeGrowsError::PARTIAL_SEND);
      }
 
      // We've most likely at least partially unblocked someone, let people know
      if (sentSomething) {
        broker->readSignal.notify_all();
 
        if (allocated) {
          broker->senderSignal.notify_all();
        }
      }
      return res;
    }

    Result<void, PubSubTrySendGrowsError> try_send(const T &message) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux, std::try_to_lock_t{});
      if (!lock.owns_lock()) {
        return error(PubSubTrySendGrowsError::COULD_NOT_LOCK);
      }
 
      if (broker->closed) {
        return error(PubSubTrySendGrowsError::CHANNEL_CLOSED);
      }
      {
        // pre-check that we won't block
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          if (cur->is_full()) {
            return error(PubSubTrySendGrowsError::NO_ROOM);
          }
        }
      }
      {
        // Now that we're good, send
        auto iter = broker->inboxes.ptr_iter();
        Queue<T> *cur;
        while (iter.next().copy_if_present(cur)) {
          ensure(cur->push(message).is_success(), "PRE CHECK FAILED");
        }
        broker->readSignal.notify_all();
      }
      return success();
    }

    Result<void, PubSubTrySendGrowsError> try_send(Allocator &alloc, const T &message) {
      ensure(broker.valid(), "NOT INITIALIZED");
      auto lock = std::unique_lock(broker->mux, std::try_to_lock_t{});
      if (!lock.owns_lock()) {
        return error(PubSubTrySendGrowsError::COULD_NOT_LOCK);
      }
 
      if (broker->closed) {
        return error(PubSubTrySendGrowsError::CHANNEL_CLOSED);
      }
 
      auto iter                                 = broker->inboxes.ptr_iter();
      bool allocated                            = false;
      bool sentSomething                        = false;
      Result<void, PubSubTrySendGrowsError> res = success();
      if (Queue<T> cur; iter.next().copy_if_present(cur)) {
        allocated |= cur.is_full();
        if (auto push = cur.push(alloc, message); push.is_error()) {
          res = error(PubSubTrySendGrowsError::ALLOCATION_FAILED);
        }
        else {
          sentSomething = true;
        }
      }
 
      if (res.is_error() && sentSomething) {
        res = error(PubSubTrySendGrowsError::PARTIAL_SEND);
      }
 
      if (sentSomething) {
        broker->readSignal.notify_all();
        if (allocated) {
          broker->senderSignal.notify_all();
        }
      }
      return res;
    }
  };
 
  static_assert(is_sender<Publisher<int>>, "PUBLISHER ISN'T A SENDER");
  static_assert(is_sender_grow<Publisher<int>>, "PUBLISHER ISN'T A GROWABLE SENDER");
  static_assert(is_sender_block<Publisher<int>>, "PUBLISHER ISN'T A SENDER");
  static_assert(is_sender_block_grow<Publisher<int>>, "PUBLISHER ISN'T A GROWABLE SENDER");
  static_assert(is_sender_before<Publisher<int>>, "PUBLISHER ISN'T A SENDER");
  static_assert(is_sender_before_grow<Publisher<int>>, "PUBLISHER ISN'T A GROWABLE SENDER");
  static_assert(is_closeable<Publisher<int>>, "PUBLISHER ISN'T CLOSEABLE");
}  // namespace mtcore::thread
 
#endif  // MTCORE_THREAD_PUB_SUB_HPP