GrabBag/VrNets/TCPClient/Src/CVrTCPClient.cpp

#include "CVrTCPClient.h"
#include <sstream>
#include <thread>
#include <chrono>
#include "VrError.h"

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <fcntl.h>
#include <sys/select.h>
#include <netinet/tcp.h>
#endif

#define  MAX_DATA_LEN 1024

#define PROTOCL_START	0x09060002
#define PROTOCL_END		0x09060003

IVrTCPClient* IVrTCPClient::CreateInstance()
{
    return new CVrTCPClient();
}

void IVrTCPClient::DestroyInstance(IVrTCPClient* pInstance)
{
    if (pInstance != nullptr) {
        delete pInstance;
        pInstance = nullptr;
    }
}

CVrTCPClient::CVrTCPClient()
	: m_nSocket(INVALID_SOCKET_VALUE)
	, m_bRecv(false)
	, m_bRecvWorking(false)
	, m_fRecvCallback(nullptr)
	, m_fLinkcCallback(nullptr)
	, m_bLink(false)
{
	_Init();
}

CVrTCPClient::~CVrTCPClient()
{
	CloseDevice();
#ifdef _WIN32
	WSACleanup();
#endif
}


int CVrTCPClient::LinkDevice(const std::string sDevIP, int nPort, bool bReLink, LinkEventFunc linkFunc, void* pParam)
{
	m_fLinkcCallback = linkFunc;
	m_pLinkParam = pParam;
	m_sIp = sDevIP;
	m_nPort = nPort;

	int nRet = _ExecLinkDev(sDevIP, nPort);

	if (SUCCESS == nRet)
	{
		m_bLink = true;
		if (m_fLinkcCallback) m_fLinkcCallback(this, true, pParam);
	}
	

	if (bReLink)
	{
		std::thread linkThread(&CVrTCPClient::_ReLinkDevThread, this);
		linkThread.detach();
		return SUCCESS;
	}
	else
	{
		return nRet;
	}
	
}

int CVrTCPClient::StartWork(TCPRecvFunc fRecvFunc, void* pParam)
{
	if (m_bRecv) return SUCCESS;

	m_bRecv = true;
	m_fRecvCallback = fRecvFunc;
	m_pWorkParam = pParam;
	std::thread recvThread(&CVrTCPClient::_RecvData, this);
	recvThread.detach();
	return SUCCESS;
}

int CVrTCPClient::CloseDevice()
{
	m_bLink = false; // 先设置连接状态为false

	if (m_bRecv)
	{
		m_bRecv = false;

		// 通知重连线程退出
		m_condRelink.notify_one();

		while (m_bRecvWorking)
		{
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}
	}

	if (m_nSocket != INVALID_SOCKET_VALUE)
	{
		closesocket_func(m_nSocket);
		m_nSocket = INVALID_SOCKET_VALUE;
	}
	return SUCCESS;
}

bool CVrTCPClient::SendData(const char* pData, const int nLen)
{
	bool bRet = true;
	int nSendLen = 0;
	do 
	{
		int nCount = 0;
		if ((nCount = send(m_nSocket, pData + nSendLen, nLen - nSendLen, 0)) < 0)
		{
			printf("send faile");
			bRet = false;
			break;
		}
		else
		{
			nSendLen += nCount;
		}
	} while (nSendLen < nLen);	
	return bRet;
}

bool CVrTCPClient::_Init()
{
#ifdef _WIN32
	WORD wVersionRequested;
	WSADATA wsaData;
	int err;
	wVersionRequested = MAKEWORD(2, 2);
	err = WSAStartup(wVersionRequested, &wsaData);
	if (err != 0)
	{
		return false;
	}
#endif
	return true;
}

void CVrTCPClient::_RecvData()
{
	m_bRecvWorking = true;
	char recvData[MAX_BUF_LEN];
	int recvLen = 0;
	while (m_bRecv)
	{
		timeval waitTime = {0, 1000};
		fd_set rfd;
		FD_ZERO(&rfd);
		FD_SET(m_nSocket, &rfd);

		int nCount = select(m_nSocket + 1, &rfd, NULL, NULL, &waitTime);
		if (nCount <= 0){
			continue;
		}

		if(!m_bLink)
		{
			continue;
		}

		memset(recvData, 0, MAX_BUF_LEN);

		//recv
		if ((recvLen = recv(m_nSocket, recvData, MAX_BUF_LEN, 0)) <= 0)
		{
			// 立即更新连接状态
			m_bLink = false;
			if (m_fLinkcCallback) m_fLinkcCallback(this, false, m_pLinkParam);

			// 通知重连线程
			m_condRelink.notify_one();
			continue;
		}
		
		if (m_fRecvCallback)
		{
			m_fRecvCallback(this, recvData, recvLen, m_pWorkParam);
		}
	} 
	m_bRecvWorking = false;
}

// 重新连接线程
void CVrTCPClient::_ReLinkDevThread()
{
	while (m_bRecv) // 改为基于接收状态而非工作状态
	{
		if (m_bLink)
		{
			std::unique_lock<std::mutex> lck(m_mutexRelink);
			m_condRelink.wait(lck);
		}

		// 如果已停止接收，退出重连线程
		if (!m_bRecv) break;

		int nLinkRet = _ExecLinkDev(m_sIp, m_nPort);

		m_bLink = ( 0 == nLinkRet );
		if (m_fLinkcCallback) m_fLinkcCallback(this, m_bLink, m_pLinkParam);

		// 如果连接失败，等待3秒再重试，避免频繁重连占用CPU
		if (!m_bLink && m_bRecv) {
			std::this_thread::sleep_for(std::chrono::seconds(3));
		}
	}
}

int CVrTCPClient::_ExecLinkDev(std::string sIP, int nPort)
{
	if (INVALID_SOCKET_VALUE != m_nSocket)
	{
		closesocket_func(m_nSocket);
	}
	m_nSocket = socket(AF_INET, SOCK_STREAM, 0);
	if (INVALID_SOCKET_VALUE == m_nSocket)
	{
		return ERR_CODE(NET_ERR_CREAT_INIT);
	}

	// 设置socket为非阻塞模式
#ifdef _WIN32
	u_long mode = 1;
	ioctlsocket(m_nSocket, FIONBIO, &mode);
#else
	int flags = fcntl(m_nSocket, F_GETFL, 0);
	fcntl(m_nSocket, F_SETFL, flags | O_NONBLOCK);
#endif

	sockaddr_in sSockAddr;
	sSockAddr.sin_family = AF_INET;
	sSockAddr.sin_port = htons(nPort);

#ifdef _WIN32
	sSockAddr.sin_addr.S_un.S_addr = inet_addr(sIP.c_str());
#else
	inet_pton(AF_INET, sIP.c_str(), &sSockAddr.sin_addr);
#endif

	int nRet = connect(m_nSocket, (sockaddr*)&sSockAddr, sizeof(sockaddr_in));
	if (nRet == SOCKET_ERROR_VALUE)
	{
#ifdef _WIN32
		int error = WSAGetLastError();
        if (error != WSAEWOULDBLOCK)
        {
            return ERR_CODE(NET_ERR_CONNECT);
        }
#else
        int error = 0;
#endif
		// 使用select等待连接完成，设置3秒超时
		fd_set writefds;
		FD_ZERO(&writefds);
		FD_SET(m_nSocket, &writefds);

		struct timeval timeout;
		timeout.tv_sec = 3;  // 3秒超时
		timeout.tv_usec = 0;

#ifdef _WIN32
		int selectRet = select(0, NULL, &writefds, NULL, &timeout); // Windows下第一个参数被忽略
#else
		int selectRet = select(m_nSocket + 1, NULL, &writefds, NULL, &timeout);
#endif
		if (selectRet <= 0)
		{
			return ERR_CODE(NET_ERR_CONNECT);
		}

        // 检查连接是否成功
#ifdef _WIN32
		int len = sizeof(error);
		if (getsockopt(m_nSocket, SOL_SOCKET, SO_ERROR, (char*)&error, &len) < 0 || error != 0)
#else
		socklen_t len = sizeof(error);
		if (getsockopt(m_nSocket, SOL_SOCKET, SO_ERROR, &error, &len) < 0 || error != 0)
#endif
		{
			return ERR_CODE(NET_ERR_CONNECT);
		}
	}

	// 恢复socket为阻塞模式
#ifdef _WIN32
	u_long blockMode = 0;
	ioctlsocket(m_nSocket, FIONBIO, &blockMode);
#else
	flags = fcntl(m_nSocket, F_GETFL, 0);
	fcntl(m_nSocket, F_SETFL, flags & ~O_NONBLOCK);
#endif

	int flag = 1;
	int ret = setsockopt(m_nSocket, IPPROTO_TCP, TCP_NODELAY, (char*)&flag, sizeof(flag));
	if (ret == -1)
	{
		printf("Couldn't setsockopt(TCP_NODELAY)\n");
	}

	return SUCCESS;
}