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              : #ifndef HCCL_HCCL_SOCKET_H
      11              : #define HCCL_HCCL_SOCKET_H
      12              : 
      13              : #include <vector>
      14              : #include <memory>
      15              : #include <atomic>
      16              : #include "hccl/hccl_types.h"
      17              : #include "hccl_common.h"
      18              : #include "hccl_ip_address.h"
      19              : #include "hccl_network_pub.h"
      20              : #include "adapter_hccp_common.h"
      21              : #include "externalinput_pub.h"
      22              : 
      23              : namespace hccl {
      24              : constexpr u32 SERVER_ROLE_SOCKET = 0;
      25              : constexpr u32 CLIENT_ROLE_SOCKET = 1;
      26              : 
      27              : constexpr u32 NIC_SOCKET_CONN_LIMIT = 1;
      28              : constexpr u32 VNIC_SOCKET_CONN_LIMIT = 1;
      29              : constexpr u32 SOCKET_BATCH_GET_LIMIT = 16; // 每次最多建立16个socket连接
      30              : constexpr u32 HOST_SOCKET_CONN_LIMIT = 16; // 用于host socket建链白名单，限制每个server的socket连接数
      31              : 
      32              : enum class HcclSocketType {
      33              :     SOCKET_NIC,
      34              :     SOCKET_HOST_NIC,
      35              :     SOCKET_VNIC,
      36              : };
      37              : 
      38              : enum class HcclSocketRole {
      39              :     SOCKET_ROLE_SERVER = 0, /* server 角色 */
      40              :     SOCKET_ROLE_CLIENT = 1, /* client 角色 */
      41              :     SOCKET_ROLE_RESERVED    /* 作为Listen Socket，或标识无需创建socket连接 */
      42              : };
      43              : 
      44              : enum class HcclSocketStatus {
      45              :     SOCKET_INIT = 0,
      46              :     SOCKET_CONNECTING = 1,
      47              :     SOCKET_OK = 2,
      48              :     SOCKET_TIMEOUT = 3,
      49              :     SOCKET_ERROR = 4,
      50              : };
      51              : 
      52              : // 如果一个Rank有多个IP里, 使用 std::vector<HcclRankLinkInfo> 描述
      53              : using HcclRankLinkInfo = struct HcclRankLinkInfoDef {
      54              :     u32 userRank;
      55              :     u32 devicePhyId;
      56              :     HcclIpAddress ip;
      57              :     u32 port;
      58              :     u32 socketsPerLink;
      59              : 
      60         1970 :     HcclRankLinkInfoDef() : userRank(), devicePhyId(), ip(), port(), socketsPerLink() {}
      61              : };
      62              : 
      63              : class HcclSocket {
      64              : public:
      65              :     explicit HcclSocket(
      66              :         const std::string& tag, HcclNetDevCtx netDevCtx, const HcclIpAddress& remoteIp, u32 remotePort,
      67              :         HcclSocketRole localRole);
      68              :     explicit HcclSocket(HcclNetDevCtx netDevCtx, u32 localPort = HCCL_INVALID_PORT);
      69              : 
      70              :     ~HcclSocket();
      71              :     HcclResult Init();
      72              :     HcclResult DeInit();
      73              :     HcclResult Listen();
      74              :     HcclResult Listen(u32 port);
      75              :     HcclResult Connect();
      76              :     void Close();
      77              :     void SetStatus(HcclSocketStatus status);
      78              :     HcclSocketStatus GetStatus();
      79              :     HcclResult Accept(const std::string& tag, std::shared_ptr<HcclSocket>& socket, u32 acceptTimeOut = 0);
      80              :     HcclResult Send(const void* data, u64 size);
      81              :     HcclResult Recv(void* recvBuf, u32 recvBufLen, u32 timeout = 0);
      82              :     HcclResult Send(const std::string& sendMsg);
      83              :     HcclResult Recv(std::string& recvMsg, u32 timeout = 0);
      84              :     HcclResult ISend(void* data, u64 size, u64& compSize);
      85              :     HcclResult IRecv(void* recvBuf, u32 recvBufLen, u64& compSize);
      86              : 
      87              :     static bool IsSupportAsync();
      88              :     HcclResult SendAsync(const void* data, u64 size, u64* sentSize, void** reqHandle);
      89              :     HcclResult RecvAsync(void* recvBuf, u64 recvBufLen, u64* receivedSize, void** reqHandle);
      90              :     HcclResult GetAsyncReqResult(void* reqHandle, HcclResult& reqResult);
      91              : 
      92              :     HcclResult AddWhiteList(std::vector<SocketWlistInfo>& wlistInfoVec);
      93              :     HcclResult DelWhiteList(std::vector<SocketWlistInfo>& wlistInfoVec);
      94              : 
      95              :     std::string GetTag() const;
      96              :     NicType GetSocketType() const;
      97              :     HcclIpAddress GetRemoteIp() const;
      98              :     u32 GetRemotePort() const;
      99              :     HcclIpAddress GetLocalIp() const;
     100              :     u32 GetLocalPort() const;
     101              :     HcclSocketRole GetLocalRole() const;
     102              :     FdHandle GetFdHandle() const;
     103              :     void SetForceClose(bool forceClose);
     104              :     HcclResult SetStopFlag(bool value);
     105              :     bool GetStopFlag();
     106              : 
     107              : private:
     108              :     HcclSocketStatus ConvertRaSocketStatus(int raStatus);
     109              :     HcclResult GetNicSocketHandle();
     110              :     HcclResult GetNicSocketHandle(
     111              :         std::map<HcclIpAddress, IpSocket>& socketMap, const HcclIpAddress& ip, SocketHandle& nicSocketHandle);
     112              :     HcclResult ListenHostNet(HcclResult& ret, std::string& errormessage, u32& port);
     113              : 
     114              :     std::string tag_;
     115              :     HcclNetDevCtx netDevCtx_;
     116              :     NICDeployment nicDeployment_;
     117              :     NicType socketType_;
     118              :     s32 localDeviceLogicId_;
     119              :     SocketHandle nicSocketHandle_{nullptr};
     120              :     s32 localDevicePhyId_;
     121              :     HcclIpAddress remoteIp_;
     122              :     u32 remotePort_;
     123              :     HcclIpAddress localIp_;
     124              :     HcclIpAddress backupIp_;
     125              :     u32 localPort_;
     126              :     HcclSocketRole localRole_;
     127              :     HcclSocketStatus status_;
     128              :     FdHandle fdHandle_;
     129              :     bool isHostUseDevNic_{false};
     130              :     bool listened_{false};
     131              :     bool forceClose_{false};
     132              :     std::atomic<bool> stopFlag_{false};
     133              :     s32 sendStatus_{0};
     134              :     s32 recvStatus_{0};
     135              : };
     136              : 
     137              : } // namespace hccl
     138              : #endif // HCCL_HCCL_SOCKET_H