taps.hpp

#pragma once
#include <thread>
 
#include <spdlog/spdlog.h>
#include <zmq.hpp>
 
#include "synapse-app-sdk/middleware/conversions.hpp"
 
#include "api/datatype.pb.h"
#include "api/tap.pb.h"
 
namespace synapse {

const std::string kDefaultTapZMQIPC = "ipc:///tmp/tap_registry";

class TapSocket {
 public:
  zmq::socket_t socket;

  synapse::TapConnection connection;
};

class ConsumerTap {
 public:
  ConsumerTap(const std::string& name, const std::string& message_type_name,
              const std::function<void(zmq::message_t message)>& callback,
              zmq::context_t& zmq_context)
      : callback_(callback), zmq_context_(zmq_context) {
    socket_ = zmq::socket_t(zmq_context_, zmq::socket_type::sub);
    socket_.set(zmq::sockopt::subscribe, "");
    socket_.bind("tcp://*:0");
    connection_.set_name(name);
    connection_.set_message_type(message_type_name);
    connection_.set_endpoint(socket_.get(zmq::sockopt::last_endpoint));
    connection_.set_tap_type(synapse::TapType::TAP_TYPE_CONSUMER);
 
    spdlog::info("Consumer tap {} <{}> created at {}", name, message_type_name,
                 connection_.endpoint());
  }

  ~ConsumerTap() { stop(); }

  synapse::TapConnection connection() const { return connection_; }

  void start() {
    running_ = true;
    read_thread_ = std::thread([this]() {
      while (running_) {
        zmq::message_t message;
        try {
          auto ret = socket_.recv(message, zmq::recv_flags::dontwait);
          if (!ret.has_value()) {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            continue;
          }
          callback_(std::move(message));
        } catch (const std::exception& e) {
          spdlog::error("Error receiving message from consumer tap {}: {}", connection_.name(),
                        e.what());
          std::this_thread::sleep_for(std::chrono::milliseconds(100));
          continue;
        }
      }
    });
  }

  void stop() {
    running_ = false;
    if (read_thread_.joinable()) {
      read_thread_.join();
    }
  }
 
 private:
  zmq::socket_t socket_;
  zmq::context_t& zmq_context_;
  std::function<void(zmq::message_t message)> callback_;

  synapse::TapConnection connection_;

  std::thread read_thread_;

  std::atomic<bool> running_{false};
};

class TapManager {
 public:
  explicit TapManager(zmq::context_t& zmq_context,
                      const std::string& registry_endpoint = kDefaultTapZMQIPC);

  ~TapManager();

  std::optional<synapse::TapConnection> add_tap(const std::string& name,
                                                const std::string& message_type_name);

  std::optional<synapse::TapConnection> add_consumer_tap(
      const std::string& name, const std::string& message_type_name,
      const std::function<void(zmq::message_t message)>& callback) {
    if (name.empty() || message_type_name.empty()) {
      return std::nullopt;
    }
 
    std::lock_guard<std::mutex> lock(consumer_tap_mutex_);
    if (consumer_taps_.count(name)) {
      return std::nullopt;
    }
    spdlog::info("Trying to add consumer tap: {} <{}>", name, message_type_name);
    auto consumer_tap =
        std::make_shared<ConsumerTap>(name, message_type_name, callback, zmq_context_);
    consumer_tap->start();
    consumer_taps_[name] = consumer_tap;
    return consumer_tap->connection();
  }

  template <typename T>
  bool send_message(const std::string& name, const T& message) {
    if (name.empty()) {
      return false;
    }
 
    // Make sure we have this tap
    if (!taps_.count(name)) {
      return false;
    }
 
    auto out_message = protobuf_to_zmq_message(message);
    const auto ret = taps_.at(name)->socket.send(out_message, zmq::send_flags::none);
 
    // Non block send
    return ret.has_value();
  }

  void stop();
 
 private:
  std::atomic<bool> running_{false};

  zmq::context_t& zmq_context_;

  std::thread tap_thread_;

  zmq::socket_t req_socket_;

  std::mutex tap_mutex_;

  std::unordered_map<std::string, std::shared_ptr<TapSocket>> taps_;
 
  std::mutex consumer_tap_mutex_;
 
  std::unordered_map<std::string, std::shared_ptr<ConsumerTap>> consumer_taps_;

  void run_tap_thread();
};
}  // namespace synapse