Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef HCCL_HCCL_SOCKET_H
      12              : #define HCCL_HCCL_SOCKET_H
      13              : 
      14              : #include <vector>
      15              : #include <memory>
      16              : #include <atomic>
      17              : #include "hccl/hccl_types.h"
      18              : #include "hccl_common.h"
      19              : #include "hccl_ip_address.h"
      20              : #include "hccl_network_pub.h"
      21              : #include "adapter_hccp_common.h"
      22              : #include "externalinput_pub.h"
      23              : 
      24              : namespace hccl {
      25              : constexpr u32 SERVER_ROLE_SOCKET = 0;
      26              : constexpr u32 CLIENT_ROLE_SOCKET = 1;
      27              : 
      28              : constexpr u32 NIC_SOCKET_CONN_LIMIT = 1;
      29              : constexpr u32 VNIC_SOCKET_CONN_LIMIT = 1;
      30              : constexpr u32 SOCKET_BATCH_GET_LIMIT = 16; // 每次最多建立16个socket连接
      31              : constexpr u32 HOST_SOCKET_CONN_LIMIT = 16; // 用于host socket建链白名单，限制每个server的socket连接数
      32              : 
      33              : enum class HcclSocketType {
      34              :     SOCKET_NIC,
      35              :     SOCKET_HOST_NIC,
      36              :     SOCKET_VNIC,
      37              : };
      38              : 
      39              : enum class HcclSocketRole {
      40              :     SOCKET_ROLE_SERVER = 0, /* server 角色 */
      41              :     SOCKET_ROLE_CLIENT = 1, /* client 角色 */
      42              :     SOCKET_ROLE_RESERVED    /* 作为Listen Socket，或标识无需创建socket连接 */
      43              : };
      44              : 
      45              : enum class HcclSocketStatus {
      46              :     SOCKET_INIT = 0,
      47              :     SOCKET_CONNECTING = 1,
      48              :     SOCKET_OK = 2,
      49              :     SOCKET_TIMEOUT = 3,
      50              :     SOCKET_ERROR = 4,
      51              : };
      52              : 
      53              : // 如果一个Rank有多个IP里, 使用 std::vector<HcclRankLinkInfo> 描述
      54              : using HcclRankLinkInfo = struct HcclRankLinkInfoDef {
      55              :     u32 userRank;
      56              :     u32 devicePhyId;
      57              :     HcclIpAddress ip;
      58              :     u32 port;
      59              :     u32 socketsPerLink;
      60              : 
      61         1970 :     HcclRankLinkInfoDef() : userRank(), devicePhyId(), ip(), port(), socketsPerLink() {}
      62              : };
      63              : 
      64              : class HcclSocket {
      65              : public:
      66              :     explicit HcclSocket(
      67              :         const std::string& tag, HcclNetDevCtx netDevCtx, const HcclIpAddress& remoteIp, u32 remotePort,
      68              :         HcclSocketRole localRole);
      69              :     explicit HcclSocket(HcclNetDevCtx netDevCtx, u32 localPort = HCCL_INVALID_PORT);
      70              : 
      71              :     ~HcclSocket();
      72              :     HcclResult Init();
      73              :     HcclResult DeInit();
      74              :     HcclResult Listen();
      75              :     HcclResult Listen(u32 port);
      76              :     HcclResult Connect();
      77              :     void Close();
      78              :     void SetStatus(HcclSocketStatus status);
      79              :     HcclSocketStatus GetStatus();
      80              :     HcclResult Accept(const std::string& tag, std::shared_ptr<HcclSocket>& socket, u32 acceptTimeOut = 0);
      81              :     HcclResult Send(const void* data, u64 size);
      82              :     HcclResult Recv(void* recvBuf, u32 recvBufLen, u32 timeout = 0);
      83              :     HcclResult Send(const std::string& sendMsg);
      84              :     HcclResult Recv(std::string& recvMsg, u32 timeout = 0);
      85              :     HcclResult ISend(void* data, u64 size, u64& compSize);
      86              :     HcclResult IRecv(void* recvBuf, u32 recvBufLen, u64& compSize);
      87              : 
      88              :     static bool IsSupportAsync();
      89              :     HcclResult SendAsync(const void* data, u64 size, u64* sentSize, void** reqHandle);
      90              :     HcclResult RecvAsync(void* recvBuf, u64 recvBufLen, u64* receivedSize, void** reqHandle);
      91              :     HcclResult GetAsyncReqResult(void* reqHandle, HcclResult& reqResult);
      92              : 
      93              :     HcclResult AddWhiteList(std::vector<SocketWlistInfo>& wlistInfoVec);
      94              :     HcclResult DelWhiteList(std::vector<SocketWlistInfo>& wlistInfoVec);
      95              : 
      96              :     std::string GetTag() const;
      97              :     NicType GetSocketType() const;
      98              :     HcclIpAddress GetRemoteIp() const;
      99              :     u32 GetRemotePort() const;
     100              :     HcclIpAddress GetLocalIp() const;
     101              :     u32 GetLocalPort() const;
     102              :     HcclSocketRole GetLocalRole() const;
     103              :     FdHandle GetFdHandle() const;
     104              :     void SetForceClose(bool forceClose);
     105              :     HcclResult SetStopFlag(bool value);
     106              :     bool GetStopFlag();
     107              : 
     108              : private:
     109              :     HcclSocketStatus ConvertRaSocketStatus(int raStatus);
     110              :     HcclResult GetNicSocketHandle();
     111              :     HcclResult GetNicSocketHandle(
     112              :         std::map<HcclIpAddress, IpSocket>& socketMap, const HcclIpAddress& ip, SocketHandle& nicSocketHandle);
     113              :     HcclResult ListenHostNet(HcclResult& ret, std::string& errormessage, u32& port);
     114              : 
     115              :     std::string tag_;
     116              :     HcclNetDevCtx netDevCtx_;
     117              :     NICDeployment nicDeployment_;
     118              :     NicType socketType_;
     119              :     s32 localDeviceLogicId_;
     120              :     SocketHandle nicSocketHandle_{nullptr};
     121              :     s32 localDevicePhyId_;
     122              :     HcclIpAddress remoteIp_;
     123              :     u32 remotePort_;
     124              :     HcclIpAddress localIp_;
     125              :     HcclIpAddress backupIp_;
     126              :     u32 localPort_;
     127              :     HcclSocketRole localRole_;
     128              :     HcclSocketStatus status_;
     129              :     FdHandle fdHandle_;
     130              :     bool isHostUseDevNic_{false};
     131              :     bool listened_{false};
     132              :     bool forceClose_{false};
     133              :     std::atomic<bool> stopFlag_{false};
     134              :     s32 sendStatus_{0};
     135              :     s32 recvStatus_{0};
     136              : };
     137              : 
     138              : } // namespace hccl
     139              : #endif // HCCL_HCCL_SOCKET_H