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