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 SHARED_JETTY_CHANNEL_POOL_H
11 : #define SHARED_JETTY_CHANNEL_POOL_H
12 :
13 : #include <cstdint>
14 : #include <cstring>
15 : #include <functional>
16 : #include <memory>
17 : #include <mutex>
18 : #include <string>
19 : #include <unordered_map>
20 : #include <vector>
21 : #include "hcomm_res_defs.h"
22 : #include "hcomm_channel.h"
23 : #include "hcomm_res.h"
24 : #include "hccl/hccl_types.h"
25 :
26 : namespace hccl {
27 :
28 : class MyRank;
29 :
30 : using EndpointDescPair = std::pair<EndpointDesc, EndpointDesc>;
31 :
32 : // 字段级 hash,规避 EndpointDesc padding 字节未初始化导致的误判。
33 : // 将参与区分的有效字段序列化到 std::string,复用标准库 std::hash<string> 完成组合,
34 : // 避免手写魔数/位运算 combine(如 0x9e3779b9)带来的可读性与稳健性问题。
35 : struct EndpointDescPairHash {
36 90 : static void AppendEndpointDesc(std::string& s, const EndpointDesc& d)
37 : {
38 90 : s.append(reinterpret_cast<const char*>(&d.protocol), sizeof(d.protocol));
39 90 : s.append(reinterpret_cast<const char*>(&d.commAddr.type), sizeof(d.commAddr.type));
40 : // commAddr.union 内 raws[36] 覆盖全部 union 存储,可安全用于 hash
41 90 : s.append(reinterpret_cast<const char*>(d.commAddr.raws), sizeof(d.commAddr.raws));
42 90 : s.append(reinterpret_cast<const char*>(&d.loc.locType), sizeof(d.loc.locType));
43 90 : s.append(reinterpret_cast<const char*>(d.loc.raws), sizeof(d.loc.raws));
44 90 : }
45 45 : std::size_t operator()(const EndpointDescPair& p) const noexcept
46 : {
47 45 : std::string buf;
48 45 : buf.reserve(sizeof(EndpointDesc) * 2);
49 45 : AppendEndpointDesc(buf, p.first);
50 45 : AppendEndpointDesc(buf, p.second);
51 45 : return std::hash<std::string>{}(buf);
52 45 : }
53 : };
54 :
55 : // 字段级比较,规避 EndpointDesc padding 字段未初始化导致的误判
56 : struct EndpointDescPairEqual {
57 29 : bool operator()(const EndpointDescPair& a, const EndpointDescPair& b) const noexcept
58 : {
59 29 : return a.first.protocol == b.first.protocol && a.first.commAddr.type == b.first.commAddr.type
60 29 : && std::memcmp(a.first.commAddr.raws, b.first.commAddr.raws, sizeof(a.first.commAddr.raws)) == 0
61 29 : && a.first.loc.locType == b.first.loc.locType
62 29 : && std::memcmp(a.first.loc.raws, b.first.loc.raws, sizeof(a.first.loc.raws)) == 0
63 29 : && a.second.protocol == b.second.protocol && a.second.commAddr.type == b.second.commAddr.type
64 29 : && std::memcmp(a.second.commAddr.raws, b.second.commAddr.raws, sizeof(a.second.commAddr.raws)) == 0
65 29 : && a.second.loc.locType == b.second.loc.locType
66 58 : && std::memcmp(a.second.loc.raws, b.second.loc.raws, sizeof(a.second.loc.raws)) == 0;
67 : }
68 : };
69 :
70 : /**
71 : * @note 职责:MyRank 级别的共享 Jetty Channel 池。
72 : * 按 myRank -> tag -> (localEndpoint, remoteEndpoint) -> [ChannelHandle] 管理复用的 Channel。
73 : * 重复调用时,若 tag->endpointPair 下 channel 数量不足则创建后再返回;若足够直接按序返回。
74 : * MyRank 析构时通过 DestroyAllByMyRank 统一清理,避免悬挂引用。
75 : * 约束:池内 Channel 生命周期由池统一管理,调用方不可单独对其调用
76 : * HcommChannelDestroy/HcclChannelRelease,否则池内会残留悬挂句柄。
77 : */
78 : class SharedJettyChannelPool {
79 : public:
80 : struct EpPairChannels {
81 : std::vector<ChannelHandle> channels;
82 : uint32_t nextReturnIdx{0};
83 : };
84 :
85 : static SharedJettyChannelPool& GetInstance();
86 :
87 : /**
88 : * @brief 获取或创建共享 Jetty Channel
89 : * @param[in] myRank 归属的 MyRank 指针(用于销毁时索引)
90 : * @param[in] tag 共享队列 tag
91 : * @param[in] epPair 源目的 endpointPair
92 : * @param[in] requestedNum 请求的 channel 数量
93 : * @param[in] createFunc 创建新 channel 的回调
94 : * @param[out] outChannels 输出的 channel 句柄数组(前 outReusedCount 个为池中复用,其余为新建)
95 : * @param[out] outReusedCount 输出从池中复用的 channel 数量(可选,nullptr 时不输出)
96 : * @return HcclResult 执行结果
97 : */
98 : HcclResult AcquireChannels(
99 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, uint32_t requestedNum,
100 : const std::function<HcclResult(uint32_t, ChannelHandle*)>& createFunc, ChannelHandle* outChannels,
101 : uint32_t* outReusedCount = nullptr);
102 :
103 : /**
104 : * @brief 销毁 MyRank 下所有共享 Jetty Channel(MyRank 析构时调用)
105 : * @param[in] myRank MyRank 指针
106 : */
107 : HcclResult DestroyAllByMyRank(MyRank* myRank);
108 :
109 : /**
110 : * @brief 检查 MyRank 是否有共享 Jetty Channel 未销毁
111 : * @param[in] myRank MyRank 指针
112 : * @return HcclResult HCCL_SUCCESS 表示可以销毁,HCCL_E_UNAVAIL 表示仍有 Channel
113 : */
114 : HcclResult CheckMyRankDestroy(MyRank* myRank);
115 :
116 : /**
117 : * @brief 从池中移除指定 channel 句柄(建链失败时清理已销毁的句柄)
118 : * @param[in] myRank 归属的 MyRank 指针
119 : * @param[in] tag 共享队列 tag
120 : * @param[in] epPair 源目的 endpointPair
121 : * @param[in] channels 要移除的 channel 句柄数组
122 : * @param[in] channelNum 数量
123 : */
124 : void RemoveChannels(
125 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, const ChannelHandle* channels,
126 : uint32_t channelNum);
127 :
128 : private:
129 5 : SharedJettyChannelPool() = default;
130 5 : ~SharedJettyChannelPool() = default;
131 : SharedJettyChannelPool(const SharedJettyChannelPool&) = delete;
132 : SharedJettyChannelPool& operator=(const SharedJettyChannelPool&) = delete;
133 :
134 : HcclResult ReturnExistingChannels(
135 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, uint32_t requestedNum,
136 : ChannelHandle* outChannels, uint32_t& returnFromExisting, uint32_t& needCreate);
137 :
138 : using EpPairMap = std::unordered_map<EndpointDescPair, EpPairChannels, EndpointDescPairHash, EndpointDescPairEqual>;
139 : using TagMap = std::unordered_map<std::string, EpPairMap>;
140 : using RankPoolIter = std::unordered_map<MyRank*, TagMap>::iterator;
141 :
142 : // 调用者须持有 mtx_。收集 myRank 下全部共享 Jetty channel 句柄。
143 : // 返回值:rankPools_ 中 myRank 对应的迭代器;未找到返回 rankPools_.end()。
144 : // allChannels 输出收集到的句柄(已 reserve 预分配)。
145 : RankPoolIter CollectMyRankChannelsLocked(MyRank* myRank, std::vector<ChannelHandle>& allChannels);
146 :
147 : std::mutex mtx_;
148 : std::unordered_map<MyRank*, TagMap> rankPools_;
149 : };
150 :
151 : } // namespace hccl
152 :
153 : #endif // SHARED_JETTY_CHANNEL_POOL_H
|