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_OPRETRY_CONNECTION_H
      12              : #define HCCL_OPRETRY_CONNECTION_H
      13              : 
      14              : #include <map>
      15              : #include <mutex>
      16              : #include <thread>
      17              : #include <memory>
      18              : #include <atomic>
      19              : #include "hccl_socket.h"
      20              : #include "hccl_network_pub.h"
      21              : #include "universal_concurrent_map.h"
      22              : #include "hccl_op_retry_pub.h"
      23              : 
      24              : namespace hccl {
      25              : const std::string OP_RETRY_CONN_SOCKET_TAG = "OP_RETRY_CONN_SOCKET_TAG";
      26              : const u32 OP_RETRY_CONN_PORT_MAX_RANGE = 15;
      27              : 
      28              : class OpRetryConnection {
      29              : public:
      30              :     OpRetryConnection();
      31              :     ~OpRetryConnection();
      32              : 
      33              :     using OpRetryConnectionPtr = std::shared_ptr<OpRetryConnection>;
      34              : 
      35              :     /* 配置是否开启该建链功能，如果不开启该功能部分接口会直接返回 */
      36              :     static void SetOpRetryConnEnable(bool enable);
      37              :     static bool IsOpRetryConnEnable();
      38              : 
      39              :     /* 创建并初始化对应全局静态资源实例*/
      40              :     static HcclResult Init(const std::string &group, u32 rankSize, const OpRetryServerInfo& serverInfo,
      41              :         const OpRetryAgentInfo& agentInfo, u32 rootRank = 0);
      42              : 
      43              :     /*
      44              :      * 供创建全局实例，以group为Key
      45              :      * forceNew表示如果已有存在则释放，创建新的实例，用户一般使用默认配置false即可
      46              :      */
      47              :     static HcclResult GetInstance(const std::string &group, OpRetryConnectionPtr &conn, bool forceNew = false);
      48              :     static HcclResult DelInstance(const std::string &group);
      49              : 
      50              :     HcclResult Init(u32 rankId, u32 rankSize, const HcclIpAddress &serverIp, u32 serverPort, s32 serverDevId,
      51              :         const HcclIpAddress &localIp, u32 rootRank = 0);
      52              :     HcclResult DeInit();
      53              : 
      54           14 :     bool IsRoot() const  /* 判断自己是否为Root节点 */
      55              :     {
      56           14 :         return rankId_ == rootRank_;
      57              :     }
      58              : 
      59              :     HcclResult GetAgentSocket(std::shared_ptr<HcclSocket> &sock);
      60              :     HcclResult GetServerSockets(std::map<u32, std::shared_ptr<HcclSocket>> &socks);
      61              : 
      62              : private:
      63            4 :     void SetGroup(const std::string &group)
      64              :     {
      65            4 :         group_ = group;
      66            4 :     }
      67              : 
      68           14 :     const std::string& GetTag()
      69              :     {
      70           14 :         if (tag_.empty()) {
      71            7 :             tag_ = OP_RETRY_CONN_SOCKET_TAG + "_" + group_ + "_" + std::to_string(serverPort_);
      72              :         }
      73              : 
      74           14 :         return tag_;
      75              :     }
      76              : 
      77              :     HcclResult InitHcclNet();
      78              :     HcclResult LoadHostWhiteList(const std::string &whiteListFile);
      79              : 
      80              :     HcclResult StartListen();
      81              :     HcclResult StopListen();
      82              :     HcclResult Accept();
      83              :     HcclResult WaitAcceptFinish();  /* 阻塞等待accept完成 */
      84              :     HcclResult RecvMetaInfo(std::shared_ptr<HcclSocket> &peerSocket);
      85              :     HcclResult SendAckInfo(std::shared_ptr<HcclSocket> &peerSocket);
      86              :     void RunAccept();               /* 线程入口，异步接收所有连接 */
      87              : 
      88              :     HcclResult Connect();
      89              :     HcclResult SendMetaInfo();
      90              :     HcclResult RecvAckInfo();
      91              : 
      92              :     HcclResult GetHostSocketWhiteList();    /* 从文件中解析白名单 */
      93              :     HcclResult AddListenSocketWhiteList();  /* 转换白名单到listen socket中 */
      94              : 
      95              :     static u32 GetServerPort();
      96              : 
      97              :     /* 公共数据结构 */
      98              :     static std::mutex lock_;
      99              :     static UniversalConcurrentMap<std::string, OpRetryConnectionPtr> *instance_;
     100              :     static bool enable_;
     101              : 
     102              :     s32 deviceLogicalID_{INVALID_INT};
     103              :     u32 devicePhysicID_{INVALID_UINT};
     104              :     u32 rankId_{INVALID_UINT};
     105              :     u32 rankSize_{0};
     106              :     u32 rootRank_{0};
     107              :     HcclIpAddress serverIp_;
     108              :     u32 serverPort_;
     109              :     HcclIpAddress localIp_;
     110              :     bool hcclNetInit_{false};
     111              :     std::string group_;
     112              :     std::string tag_;
     113              : 
     114              :     /* Server侧的数据结构 */
     115              :     std::shared_ptr<HcclSocket> listenSocket_{nullptr};
     116              :     std::map<u32, std::shared_ptr<HcclSocket>> connectionSockets_;
     117              :     HcclNetDevCtx serverNetCtx_{nullptr};
     118              :     bool enableWhitelist_{false};
     119              :     std::vector<HcclIpAddress> whitelist_;
     120              :     std::vector<SocketWlistInfo> wlistInfosVec_;
     121              :     std::atomic<bool> acceptFinished_{false};
     122              :     std::atomic<bool> backgroudThreadStop_{false};
     123              :     std::thread backgroudThread_;
     124              : 
     125              :     /* Client侧的数据结构 */
     126              :     std::shared_ptr<HcclSocket> socket_{nullptr};
     127              :     HcclNetDevCtx clientNetCtx_{nullptr};
     128              : };
     129              : }
     130              : 
     131              : #endif