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         1961 :     HcclRankLinkInfoDef () : userRank(), devicePhyId(), ip(), port(), socketsPerLink()
      61         1961 :     {}
      62              : };
      63              : 
      64              : class HcclSocket {
      65              : public:
      66              :     explicit HcclSocket(const std::string &tag, HcclNetDevCtx netDevCtx,
      67              :         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              : private:
     108              :     HcclSocketStatus ConvertRaSocketStatus(int raStatus);
     109              :     HcclResult GetNicSocketHandle();
     110              :     HcclResult GetNicSocketHandle(std::map<HcclIpAddress, IpSocket> &socketMap,
     111              :         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