GrabBag/Module/ZeroMQServer/Src/CVrZeroMQServer.cpp

#include "CVrZeroMQServer.h"
#include <iostream>
#include <thread>
#include <sstream>
#include <string>

#include "VrError.h"
#include "IVrUtils.h"
#include "VrStringUtils.h"
#include "VrLog.h"

bool IVrZeroMQServer::CreateObject(IVrZeroMQServer** ppMQServer)
{
	if (ppMQServer == nullptr) {
		return false;
	}
	
	try {
		*ppMQServer = new CVrZeroMQServer();
		return true;
	}
	catch (const std::exception& e) {
		std::cerr << "Failed to create ZeroMQ server: " << e.what() << std::endl;
		return false;
	}
}

CVrZeroMQServer::CVrZeroMQServer()
	: m_funMQEvent(nullptr)
	, m_funRecvAck(nullptr)
	, m_bMonitor(false)
	, m_bMonitorWorking(false)
	, m_bInitialized(false)
{
	m_pContext	= new zmq::context_t;
    m_pRouter	= new zmq::socket_t(*m_pContext, ZMQ_ROUTER);
    m_pWorkers	= new zmq::socket_t(*m_pContext, ZMQ_DEALER);

    m_pSocket 	= new zmq::socket_t(*m_pContext, ZMQ_REP);
}

CVrZeroMQServer::~CVrZeroMQServer()
{
	UnInit();

	delete m_pSocket;
	delete m_pWorkers;
	delete m_pRouter;
	delete m_pContext;
}

/// 启动服务
int CVrZeroMQServer::Init(int nPort, FunMQEvent funMQEvent, FunServerRecvAndAck funRecvAck, bool bMultitask)
{
	m_funMQEvent = funMQEvent;
	m_funRecvAck = funRecvAck;

	std::string transferAddr = "tcp://*:";
	transferAddr.append(CVrStringUtils::Int2Str(nPort));
    if (!bMultitask) {
		try {
			m_pSocket->bind(transferAddr);
				LOG_DEBUG("[ZMQ-REP] bind OK: %s\n", transferAddr.c_str());
		} catch (const std::exception& e) {
			LOG_ERROR("[ZMQ-REP] bind FAIL: %s\n", e.what());
			return ERR_CODE(NET_ERR_CREAT_BIND);
		} 
        m_bRecv = true;
        m_recvThread = std::thread(&CVrZeroMQServer::_WorkTask, this, m_pSocket);
        m_recvThread.detach();
		LOG_DEBUG("[ZMQ-REP] recv thread started, port=%d\n", nPort);
    }
    else {
        try {
            // Frontend ROUTER binds to TCP port for external clients
            m_pRouter->bind(transferAddr);
        } catch (const std::exception& e) {
            LOG_ERROR("Failed to bind ROUTER: %s", e.what());
            return ERR_CODE(NET_ERR_CREAT_BIND);
        }

        // Backend DEALER binds to inproc for worker pool
        m_pWorkers->bind("inproc://workers");

        // Spawn worker tasks based on hardware concurrency (min 2)
        unsigned int nWorkers = 5;
        for (unsigned int i = 0; i < nWorkers; ++i) {
            auto* pTask = new CVrZeroMQTask();
            pTask->Init(m_pContext, funMQEvent, funRecvAck);
            m_vecTask.push_back(pTask);
        }

        // Start proxy thread to route between ROUTER and DEALER
        m_workerThread = std::thread(&CVrZeroMQServer::_WorkerThread, this);
        m_workerThread.detach();
    }

	m_bMonitor = true;
    // 启动监控线程（监控 ROUTER 套接字，单线程模式也可复用）
    m_monitorThread = std::thread(&CVrZeroMQServer::_MonitorEvent, this, m_pRouter);
    m_monitorThread.detach();
    return SUCCESS;
}

/// 销毁服务
int CVrZeroMQServer::UnInit()
{
    m_bMonitor = false;

    m_bRecv = false;

    // 关闭单线程模式下的 socket 以唤醒阻塞的 recv
    try {
        if (m_pSocket) m_pSocket->close();
    } catch (...) {}

    // 等待接收线程/代理线程/监控线程结束
    if (m_recvThread.joinable()) m_recvThread.join();
    if (m_workerThread.joinable()) m_workerThread.join();
    if (m_monitorThread.joinable()) m_monitorThread.join();

	for (auto& pTask : m_vecTask)
	{
		pTask->UnInit();
		delete pTask;
	}
	m_vecTask.clear();

    try { if (m_pWorkers) m_pWorkers->close(); } catch (...) {}
    try { if (m_pRouter)  m_pRouter->close();  } catch (...) {}
    try { if (m_pContext) m_pContext->close(); } catch (...) {}

	return SUCCESS;
}

void CVrZeroMQServer::_WorkerThread()
{
    // 使用 zmq proxy 在线程中转发 ROUTER<->DEALER，避免阻塞主线程
    try {
        zmq::proxy(*m_pRouter, *m_pWorkers);
    } catch (const std::exception& e) {
        LOG_DEBUG("ZeroMQ proxy terminated: %s", e.what());
    }
}

void CVrZeroMQServer::on_event_handshake_succeeded(const zmq_event_t& event_, const char* addr_)
{
	LOG_DEBUG("on_event_handshake_succeeded \n");
	if (m_funMQEvent) {
		m_funMQEvent(MQ_EVENT_NEWLINK);
	}
}

void CVrZeroMQServer::on_event_disconnected(const zmq_event_t& event_, const char* addr_)
{
	LOG_DEBUG("on_event_disconnected \n");
	if (m_funMQEvent) {
		m_funMQEvent(MQ_EVENT_DELINK);
	}
}

void CVrZeroMQServer::_WorkTask(zmq::socket_t *socket)
{
    if (!socket) return;

    m_bRecvWorking = true;
    try {
        while (m_bRecv) {
            zmq::message_t msg;
            zmq::recv_result_t ret = socket->recv(msg, zmq::recv_flags::none);
            if (!m_bRecv) break;
            if (ret.has_value()) {
                LOG_DEBUG("[ZMQ-REP] _WorkTask recv: %.*s\n", (int)msg.size(), static_cast<char*>(msg.data()));
                if (m_funRecvAck) {
                    std::string strResponse = m_funRecvAck(static_cast<char*>(msg.data()), msg.size());
                    zmq::message_t response(strResponse.c_str(), strResponse.size());
                    socket->send(response, zmq::send_flags::none);
                    LOG_DEBUG("[ZMQ-REP] _WorkTask send resp bytes=%zu\n", strResponse.size());
                } else {
                    socket->send(msg, zmq::send_flags::none);
                }
            }
        }
    } catch (const std::exception& e) {
        LOG_DEBUG("Server recv loop terminated: %s", e.what());
    }
    m_bRecvWorking = false;
}

void CVrZeroMQServer::_MonitorEvent(zmq::socket_t* socket)
{
	// 设置监控
	try {
		init(*socket, "inproc://monitor", ZMQ_EVENT_ALL);
		m_bMonitorWorking = true;
		while (m_bMonitor)
		{
			check_event(100);
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}
		m_bMonitorWorking = false;
	} catch (const std::exception& e) {
		LOG_DEBUG("Monitor exception: %s", e.what());
		m_bMonitorWorking = false;
	}
}