ZMQParent.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2018 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Nils Braun                                               *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
#pragma once
 
#include <framework/logging/Logger.h>
 
#include <zmq.hpp>
#include <string>
#include <memory>
#include <thread>
#include <chrono>
 
namespace Belle2 {
  class ZMQParent {
  public:
    ~ZMQParent();
 
    template<class AZMQMessage>
    static void send(std::unique_ptr<zmq::socket_t>& socket, AZMQMessage message);
 
    static unsigned int poll(const std::vector<zmq::socket_t*>& socketList, int timeout);
 
    static std::string createIdentity(unsigned int pid = 0);
 
    void terminate();
 
    template<int AZMQType>
    std::unique_ptr<zmq::socket_t> createSocket(const std::string& socketAddress, bool bind);
 
    template<int AZMQType>
    std::unique_ptr<zmq::socket_t> createSocket(const std::string& socketAddress);
 
    void reset();
 
  private:
    std::unique_ptr<zmq::context_t> m_context;
 
    void initialize();
  };
 
  template<int AZMQType>
  std::unique_ptr<zmq::socket_t> ZMQParent::createSocket(const std::string& socketAddress, bool bind)
  {
    initialize();
 
    try {
      // TODO: options to test additionally in the future:
      // * ZMQ_IMMEDIATE
      // * ZMQ_PROBE_ROUTER
      // * ZMQ_ROUTER_MANDATORY
      // * ZMQ_ROUTER_NOTIFY
      auto socket = std::make_unique<zmq::socket_t>(*m_context, AZMQType);
 
      if (AZMQType == ZMQ_DEALER) {
        auto identity = createIdentity();
        socket->setsockopt(ZMQ_IDENTITY, identity.c_str(), identity.length());
      }
 
      // Linger means: close the socket immediately if requested, do not wait until all messages have been sent.
      // This is needed because if we want to ABORT, we usually know what we are doing (and want to do it now).
      socket->setsockopt(ZMQ_LINGER, 0);
      if (bind) {
        socket->bind(socketAddress.c_str());
      } else {
        socket->connect(socketAddress.c_str());
      }
 
      // Wait a bit to give the framework time to initialize the socket
      std::this_thread::sleep_for(std::chrono::milliseconds(100));
 
      return socket;
    } catch (zmq::error_t& error) {
      B2FATAL("Creating the ZMQ socket for address " << socketAddress << " failed because of: " << error.what());
    }
  }
 
  template<int AZMQType>
  std::unique_ptr<zmq::socket_t> ZMQParent::createSocket(const std::string& socketAddress)
  {
    // We only check for "*" in the full address for now. This is fine, as neither a valid hostname nor a valid
    // port number can ever contain a "*".
    if (socketAddress.find("*") != std::string::npos) {
      return createSocket<AZMQType>(socketAddress, true);
    } else {
      return createSocket<AZMQType>(socketAddress, false);
    }
  }
 
  template<class AZMQMessage>
  void ZMQParent::send(std::unique_ptr<zmq::socket_t>& socket, AZMQMessage message)
  {
    AZMQMessage::element_type::toSocket(std::move(message), socket);
  }
 
} // namespace Belle2