Jungfraujoch/common/ZMQWrappers.h

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only


#pragma once

#include <vector>
#include <thread>
#include <mutex>
#include <zmq.h>

#include "JFJochException.h"

class ZMQContext {
    void *context;
public:
    ZMQContext();
    ZMQContext& NumThreads(int32_t threads);
    ~ZMQContext();
    void *GetContext() const;
};

enum class ZMQSocketType : int {Push = ZMQ_PUSH, Pull = ZMQ_PULL, Req = ZMQ_REQ, Rep = ZMQ_REP,
                                Pub = ZMQ_PUB, Sub = ZMQ_SUB};

class ZMQMessage {
    zmq_msg_t msg;
public:
    ZMQMessage();
    ~ZMQMessage();
    ZMQMessage(ZMQMessage &other) = delete;
    ZMQMessage& operator=(ZMQMessage &other) = delete;
    zmq_msg_t *GetMsg();
    size_t size();
    const uint8_t *data();
};

class ZMQSocket {
    std::mutex m;
    ZMQContext context;
    ZMQSocketType socket_type;
    void *socket;
    void SetSocketOption(int32_t option_name, int32_t value);
public:
    ZMQSocket(ZMQSocket &socket) = delete;
    const ZMQSocket& operator=(ZMQSocket &socket) = delete;
    explicit ZMQSocket(ZMQSocketType socket_type);
     ~ZMQSocket();
    void Connect(const std::string& addr);
    void Disconnect(const std::string& addr);
    void Bind(const std::string& addr);
    ZMQSocket &NoReceiveTimeout();
    ZMQSocket &ReceiveTimeout(std::chrono::milliseconds input);

    ZMQSocket &NoSendTimeout();
    ZMQSocket &SendTimeout(std::chrono::milliseconds input);

    ZMQSocket &Subscribe(const std::string &topic);
    ZMQSocket &SubscribeAll();
    ZMQSocket &NoLinger();
    ZMQSocket &Conflate(bool input);
    ZMQSocket &SendBufferSize(int32_t bytes);
    ZMQSocket &ReceiverBufferSize(int32_t bytes);

    bool Receive(ZMQMessage& msg, bool blocking = true);

    void Send();
    bool Send(const void *buf, size_t buf_size, bool blocking = true, bool multipart = false);
    bool SendZeroCopy(void *data, size_t size,void (*callback)(void *, void *), void *hint, bool blocking = true);
    template <class T> bool Send(const std::vector<T> &buf) {
        return Send(buf.data(), buf.size() * sizeof(T));
    }
    void Send(const int32_t &value);
    bool Send(const std::string &s, bool blocking = true, bool multipart = false);
    void Send(zmq_msg_t *msg);
    ZMQSocket &SendWaterMark(int32_t msgs);
    ZMQSocket &ReceiveWaterMark(int32_t msgs);
    ZMQSocket &EnableKeepAlive();
    std::string GetEndpointName();
};