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() { return bodyLen != 0 && headerSended == headerLen && bodySended == bodyLen; }
      46              :     };
      47              : 
      48              :     struct FdContext {
      49              :         std::shared_ptr<Socket> socket;
      50              :         SendState txState;
      51              :     };
      52              : 
      53              :     using WorkerTask = std::function<void(void)>;
      54              : 
      55              : public:
      56              :     static constexpr u32 DEFAULT_THREAD_NUM = 1;
      57              :     static constexpr u32 MAX_THREAD_NUM = 4;
      58              :     static constexpr s32 INVALID_EPOLL_EVENT_FD = -1;
      59              :     static constexpr s32 EPOLL_TIMEOUT_MS = 100;    // 100ms
      60              :     static constexpr s32 LAST_EPOLL_TIMEOUT_MS = 5; // 5ms
      61              :     static constexpr s32 RANK_CAPACITY_PER_THREAD = 512;
      62              : 
      63           20 :     explicit RankInfoDispather(RankInfoDetectService* rankInfoDetectServer, u32 threadNum = DEFAULT_THREAD_NUM)
      64           20 :         : rankInfoDetectServer_(rankInfoDetectServer),
      65           20 :           threadNum_(threadNum)
      66           20 :     {}
      67              :     ~RankInfoDispather();
      68              : 
      69              :     void BroadcastRankTable(
      70              :         const std::unordered_map<std::string, std::shared_ptr<Socket>>& connectSockets,
      71              :         const RankTableInfo& clusterInfo, const std::string& failedAgentIdList, u32 step);
      72              : 
      73              : private:
      74              :     void InitWorkerThread();
      75              :     void WorkerWait(int workId);
      76              :     void WakeWoker();
      77              :     void RunWorkerThread(int workId);
      78              :     bool GetTask(WorkerTask& workTask);
      79              :     void PrepareResource(
      80              :         const std::unordered_map<std::string, std::shared_ptr<Socket>> connectSockets, const RankTableInfo& clusterInfo,
      81              :         const std::string& failedAgentIdList, u32 step);
      82              :     void SendOnce();
      83              :     void ProcessOneSendEvent(int epollFd, FdHandle& fdHandle);
      84              :     void ProcessSend();
      85              :     void CleanResource();
      86              :     void CloseEpollFd();
      87              : 
      88              :     RankInfoDetectService* rankInfoDetectServer_;
      89              :     u32 threadNum_{DEFAULT_THREAD_NUM};
      90              :     u32 rankNum_{0};
      91              :     std::vector<std::thread> workerThreads_;
      92              :     std::queue<WorkerTask> taskQueue_;
      93              :     std::mutex taskQueueMutex_;
      94              : 
      95              :     std::unordered_map<FdHandle, FdContext> fdHandleToFdContextMap_;
      96              :     std::mutex fdHandleMapMutex_;
      97              :     s32 epollFds_ = INVALID_EPOLL_EVENT_FD;
      98              :     std::atomic<u32> sendDoneCount_{0};
      99              : 
     100              :     std::vector<char> rankTableMsg_;
     101              : 
     102              :     std::mutex wakeMutex_;
     103              :     std::atomic<bool> epollCreate_{false};
     104              :     std::atomic<bool> ready_{false};
     105              :     std::atomic<bool> stop_{false};
     106              :     std::condition_variable wakeManager_;
     107              : };
     108              : } // namespace Hccl
     109              : 
     110              : #endif /* TOPOINFO_EXCHANGE_DISPATCHER_H */