Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 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 TOPOINFO_EXCHANGE_DISPATCHER_H
      12              : #define TOPOINFO_EXCHANGE_DISPATCHER_H
      13              :  
      14              : #include <map>
      15              : #include <atomic>
      16              : #include <vector>
      17              : #include <mutex>
      18              : #include <thread>
      19              : #include <climits>
      20              : #include <condition_variable>
      21              : #include <queue>
      22              : 
      23              : #include "types.h"
      24              : #include "socket.h"
      25              : #include "ip_address.h"
      26              : #include "rank_info_detect_service.h"
      27              : 
      28              : namespace Hccl {
      29              : 
      30              : class RankInfoDispather {
      31              : public:
      32              :     struct SendState {
      33              :         u32 rankId;
      34              :         u64 header;
      35              :         size_t headerLen       = sizeof(u64); // the header need to send
      36              :         size_t headerSended    = 0;           // the header have sended length
      37              :         size_t bodyLen         = 0;           // the whole data length
      38              :         size_t bodySended      = 0;           // the data have sended
      39              :         void *data;                           // data pointer
      40              :  
      41              :         bool Send(std::shared_ptr<Socket> socket);
      42              :         bool SendHeader(std::shared_ptr<Socket> socket);
      43              :         bool SendBody(std::shared_ptr<Socket> socket);
      44              :         bool SendHelper(std::shared_ptr<Socket> socket, void *buf, size_t dataLen, size_t &sendedLen);
      45            0 :         bool IsOk()
      46              :         {
      47            0 :             return bodyLen != 0 && headerSended == headerLen && bodySended == bodyLen;
      48              :         }
      49              :     };
      50              : 
      51              :     struct FdContext {
      52              :         std::shared_ptr<Socket> socket;
      53              :         SendState txState;
      54              :     };
      55              : 
      56              :     using WorkerTask = std::function<void(void)>;
      57              :  
      58              : public:
      59              :     static constexpr u32 DEFAULT_THREAD_NUM = 1;
      60              :     static constexpr u32 MAX_THREAD_NUM = 4;
      61              :     static constexpr s32 INVALID_EPOLL_EVENT_FD = -1;
      62              :     static constexpr s32 EPOLL_TIMEOUT_MS = 100; // 100ms
      63              :     static constexpr s32 LAST_EPOLL_TIMEOUT_MS = 5; // 5ms
      64              :     static constexpr s32 RANK_CAPACITY_PER_THREAD = 512;
      65              :  
      66           20 :     explicit RankInfoDispather(RankInfoDetectService *rankInfoDetectServer, u32 threadNum = DEFAULT_THREAD_NUM)
      67           20 :         : rankInfoDetectServer_(rankInfoDetectServer), threadNum_(threadNum)
      68              :     {
      69           20 :     }
      70              :     ~RankInfoDispather();
      71              :  
      72              :     void BroadcastRankTable(const std::unordered_map<std::string, std::shared_ptr<Socket>> &connectSockets,
      73              :         const RankTableInfo &clusterInfo, const std::string &failedAgentIdList, u32 step);
      74              :  
      75              : private:
      76              :     void InitWorkerThread();
      77              :     void WorkerWait(int workId);
      78              :     void WakeWoker();
      79              :     void RunWorkerThread(int workId);
      80              :     bool GetTask(WorkerTask &workTask);
      81              :     void PrepareResource(const std::unordered_map<std::string, std::shared_ptr<Socket>> connectSockets,
      82              :          const RankTableInfo &clusterInfo, const std::string &failedAgentIdList, u32 step);
      83              :     void SendOnce();
      84              :     void ProcessOneSendEvent(int epollFd, FdHandle &fdHandle);
      85              :     void ProcessSend();
      86              :     void CleanResource();
      87              :     void CloseEpollFd();
      88              : 
      89              :     RankInfoDetectService *rankInfoDetectServer_;
      90              :     u32 threadNum_{DEFAULT_THREAD_NUM};
      91              :     u32 rankNum_{0};
      92              :     std::vector<std::thread> workerThreads_;
      93              :     std::queue<WorkerTask> taskQueue_;
      94              :     std::mutex taskQueueMutex_;
      95              : 
      96              :     std::unordered_map<FdHandle, FdContext> fdHandleToFdContextMap_;
      97              :     std::mutex fdHandleMapMutex_;
      98              :     s32 epollFds_ = INVALID_EPOLL_EVENT_FD;
      99              :     std::atomic<u32> sendDoneCount_{0};
     100              : 
     101              :     std::vector<char> rankTableMsg_;
     102              : 
     103              :     std::mutex wakeMutex_;
     104              :     std::atomic<bool> epollCreate_{false};
     105              :     std::atomic<bool> ready_{false};
     106              :     std::atomic<bool> stop_{false};
     107              :     std::condition_variable wakeManager_;
     108              : };
     109              : } // namespace Hccl
     110              : 
     111              : #endif /* TOPOINFO_EXCHANGE_DISPATCHER_H */