优化后的TCP client 增加心跳包

优化后的TCP client 增加心跳包

#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <thread>
#include <chrono>

#pragma comment(lib, "ws2_32.lib")

#define SERVER_IP "127.0.0.1"
#define SERVER_PORT 8080
#define RECV_BUF_SIZE 1024
#define HEARTBEAT_INTERVAL 5000  // 心跳包发送间隔，单位：毫秒
#define HEARTBEAT_TIMEOUT 10000  // 心跳包超时时间，单位：毫秒

class TCPClient {
public:
    TCPClient();
    ~TCPClient();

    bool connectToServer();
    void disconnect();
    int sendData(const char* data, int len);
    int receiveData(char* buffer, int bufLen);

private:
    SOCKET m_socket;
    sockaddr_in m_serverAddr;
    
public:
    bool m_bHeartBeat;//是否启动心跳包
    bool m_connected;
    // 心跳包相关函数
    static DWORD WINAPI ThreadstartHeartbeat(LPVOID lpParam);
    int startHeartbeat();
    void sendHeartbeat();
    bool checkHeartbeatResponse();

    // 重连相关函数
    bool reconnect();
};

TCPClient::TCPClient() : m_socket(INVALID_SOCKET), 
m_connected(false), m_bHeartBeat(false)
{
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) 
    {
        std::cerr << "WSAStartup failed." << std::endl;
        // 可以在这里抛出异常或者采取其他错误处理方式
    }

    m_serverAddr.sin_family = AF_INET;
    m_serverAddr.sin_port = htons(SERVER_PORT);
    if (inet_pton(AF_INET, SERVER_IP, &(m_serverAddr.sin_addr)) <= 0) 
    {
        std::cerr << "inet_pton error." << std::endl;
        // 可以在这里添加更合适的错误处理逻辑
    }
}

TCPClient::~TCPClient()
{
    disconnect();
    WSACleanup();
}

bool TCPClient::connectToServer()
{
    m_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (m_socket == INVALID_SOCKET) {
        std::cerr << "Socket creation failed." << WSAGetLastError() << std::endl;
        return false;
    }

    int ret = connect(m_socket, (struct sockaddr*)&m_serverAddr, sizeof(m_serverAddr));
    if (ret == SOCKET_ERROR) {
        int errCode = WSAGetLastError();
        if (errCode == WSAECONNREFUSED) {
            std::cerr << "Connection refused. Will attempt to reconnect." << std::endl;
            // 尝试重连
            return reconnect();
        }
        else {
            std::cerr << "Connect failed. Error code: " << errCode << std::endl;
            closesocket(m_socket);
            m_socket = INVALID_SOCKET;
            return false;
        }
    }

    m_connected = true;
    std::cout << "Connected to server." << std::endl;

    // 连接成功后启动心跳包线程
    if (m_bHeartBeat == true)
    {
        HANDLE heartBeatThreadHandle = CreateThread(NULL, 0, ThreadstartHeartbeat, this, 0, NULL);
        if (heartBeatThreadHandle == NULL)
        {
            std::cerr << "Create heartBeat thread failed: " << GetLastError() << std::endl;
        }
        else
        {
            std::cerr << "Create heartBeat thread success: " << heartBeatThreadHandle << std::endl;
        }
    }
    return true;
}

void TCPClient::disconnect() 
{
    if (m_connected) {
        closesocket(m_socket);
        m_connected = false;
        std::cout << "Disconnected from server." << std::endl;
    }
}

int TCPClient::sendData(const char* data, int len)
{
    if (!m_connected)
    {
        std::cerr << "Not connected. Cannot send data." << std::endl;
        return SOCKET_ERROR;
    }

    int ret = send(m_socket, data, len, 0);
    if (ret == SOCKET_ERROR)
    {
        std::cerr << "Send failed. Error code: " << WSAGetLastError() << std::endl;
        // 检查错误是否是连接断开，如果是则尝试重连
        int errCode = WSAGetLastError();
        if (errCode == WSAECONNRESET || errCode == WSAENETRESET)
        {
            std::cerr << "Connection reset. Will attempt to reconnect." << std::endl;
            if (reconnect()) {
                // 重连成功后重新发送数据
                return sendData(data, len);
            }
        }
    }
    return ret;
}

int TCPClient::receiveData(char* buffer, int bufLen)
{
    if (!m_connected)
    {
        std::cerr << "Not connected. Cannot receive data." << std::endl;
        return SOCKET_ERROR;
    }

    int ret = recv(m_socket, buffer, bufLen, 0);
    if (ret == SOCKET_ERROR)
    {
        std::cerr << "Receive failed. Error code: " << WSAGetLastError() << std::endl;
        // 检查错误是否是连接断开，如果是则尝试重连
        int errCode = WSAGetLastError();
        if (errCode == WSAECONNRESET || errCode == WSAENETRESET) {
            std::cerr << "Connection reset. Will attempt to reconnect." << std::endl;
            if (reconnect()) {
                // 重连成功后重新接收数据
                return receiveData(buffer, bufLen);
            }
        }
    }
    return ret;
}

int TCPClient::startHeartbeat()
{
    while (m_connected) 
    {
        sendHeartbeat();
        if (!checkHeartbeatResponse()) 
        {
            std::cerr << "Heartbeat timeout. Will attempt to reconnect." << std::endl;
            if (reconnect())
            {
                continue;
            }
            else
            {
                // 重连失败，退出心跳包线程
                std::cout << "Heartbeat timeout. exit." << std::endl;
                m_connected = false;
                break;
            }
        }
       Sleep(HEARTBEAT_INTERVAL);
    }
    return 0;
}

// 心跳包线程函数
DWORD WINAPI TCPClient::ThreadstartHeartbeat(LPVOID lpParam)
{
    TCPClient* t_Clienet = static_cast<TCPClient*>(lpParam);
    t_Clienet->startHeartbeat();
    return 0;
}

// 发送心跳包
void TCPClient::sendHeartbeat() 
{
    const char* heartbeatData = "HEARTBEAT_CLIENT";
    int ret = sendData(heartbeatData, strlen(heartbeatData));
    if (ret > 0)
    {
        std::cout << "sent heart beat: " << heartbeatData << std::endl;
    }
    else
    {
        std::cout << "sent heart beat: " << heartbeatData << std::endl;
    }
}

// 检查心跳包响应
bool TCPClient::checkHeartbeatResponse()
{
    char buffer[RECV_BUF_SIZE] = {0};
    int ret = receiveData(buffer, RECV_BUF_SIZE);
    if (ret > 0)
    {
        std::cout << "heart beat response: " << buffer << std::endl;
        return 1;
        buffer[ret] = '\0';
        if (strcmp(buffer, "HEARTBEAT_ACK") == 0) 
        {
            return true;
        }
    }
    return false;
}

bool TCPClient::reconnect() 
{
    // 关闭当前无效的套接字
    if (m_socket != INVALID_SOCKET) {
        closesocket(m_socket);
        m_socket = INVALID_SOCKET;
    }

    // 等待一段时间后尝试重连
    Sleep(5000);  // 这里等待 5 秒，可以根据实际情况调整

    // 重新创建套接字并连接
    m_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (m_socket == INVALID_SOCKET)
    {
        std::cerr << "Socket creation failed during reconnect." << WSAGetLastError() << std::endl;
        return false;
    }

    int ret = connect(m_socket, (struct sockaddr*)&m_serverAddr, sizeof(m_serverAddr));
    if (ret == SOCKET_ERROR) {
        int errCode = WSAGetLastError();
        if (errCode == WSAECONNREFUSED) {
            std::cerr << "Reconnection refused. Will try again later." << std::endl;
            // 可以在这里设置更复杂的重连策略，比如增加等待时间后再次重连
            return false;
        }
        else {
            std::cerr << "Reconnect failed. Error code: " << errCode << std::endl;
            closesocket(m_socket);
            m_socket = INVALID_SOCKET;
            return false;
        }
    }

    m_connected = true;
    std::cout << "Reconnected to server." << std::endl;

    if (m_bHeartBeat == true)
    {
        std::thread heartbeatThread(&TCPClient::startHeartbeat, this);
        heartbeatThread.detach();
    }

    if (m_bHeartBeat == true)
    {
        HANDLE heartBeatThreadHandle = CreateThread(NULL, 0,ThreadstartHeartbeat, this, 0, NULL);
        if (heartBeatThreadHandle == NULL)
        {
            std::cerr << "Create heartBeat thread failed: " << GetLastError() << std::endl;
        }
        else
        {
            std::cerr << "Create heartBeat thread success: " << heartBeatThreadHandle << std::endl;
        }
    }

    return true;
}

int main()
{
    TCPClient client;
    client.m_bHeartBeat = false;
    static int m = 0;
    if (client.connectToServer()) 
    {
        while (client.m_connected)
        {
            Sleep(1000);
            char sendBuf[1024] = {0}; 
            sprintf(sendBuf, "Hello, server  %d", m++);
            client.sendData(sendBuf, strlen(sendBuf));

            char recvBuf[RECV_BUF_SIZE];
            int recvLen = client.receiveData(recvBuf, RECV_BUF_SIZE);
            if (recvLen > 0)
            {
                recvBuf[recvLen] = '\0';
                std::cout << "Received from server: " << recvBuf << std::endl;
            }
        }
    }
    else
    {
        std::cout << "connect to server failed" << std::endl;
    }

    return 0;
}