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              : #include "shared_jetty_mgr.h"
      11              : #include "log.h"
      12              : 
      13              : namespace hcomm {
      14              : 
      15          135 : SharedJettyMgr& SharedJettyMgr::GetInstance()
      16              : {
      17          135 :     static SharedJettyMgr instance;
      18          135 :     return instance;
      19              : }
      20              : 
      21              : HcclResult
      22           13 : SharedJettyMgr::RegisterChannels(EndpointHandle endpointHandle, const ChannelHandle* channels, uint32_t channelNum)
      23              : {
      24           13 :     if (endpointHandle == nullptr || channels == nullptr || channelNum == 0) {
      25            3 :         HCCL_ERROR(
      26              :             "[%s] invalid params, endpointHandle[%p], channels[%p], channelNum[%u].", __func__, endpointHandle,
      27              :             channels, channelNum);
      28            3 :         return HCCL_E_PARA;
      29              :     }
      30              : 
      31           10 :     std::lock_guard<std::mutex> lock(mtx_);
      32           10 :     auto& ctx = contexts_[endpointHandle];
      33           26 :     for (uint32_t i = 0; i < channelNum; ++i) {
      34           16 :         ctx.channelHandles.insert(channels[i]);
      35              :     }
      36              :     // 以集合实际大小为准：insert 对重复句柄为 no-op，若用 += channelNum 会导致
      37              :     // channelCount > channelHandles.size()，UnregisterChannels 永远无法将 count 减到 0，
      38              :     // CheckEndpointDestroy 永久阻塞 endpoint 销毁。
      39           10 :     ctx.channelCount = static_cast<uint32_t>(ctx.channelHandles.size());
      40           10 :     HCCL_INFO(
      41              :         "[%s] registered %u channels for endpointHandle[%p], total channelCount[%u].", __func__, channelNum,
      42              :         endpointHandle, ctx.channelCount);
      43           10 :     return HCCL_SUCCESS;
      44           10 : }
      45              : 
      46           25 : HcclResult SharedJettyMgr::UnregisterChannels(const ChannelHandle* channels, uint32_t channelNum)
      47              : {
      48           25 :     if (channels == nullptr || channelNum == 0) {
      49            1 :         return HCCL_SUCCESS;
      50              :     }
      51              : 
      52           24 :     std::lock_guard<std::mutex> lock(mtx_);
      53           50 :     for (uint32_t i = 0; i < channelNum; ++i) {
      54           26 :         for (auto it = contexts_.begin(); it != contexts_.end(); ++it) {
      55            5 :             auto handleIt = it->second.channelHandles.find(channels[i]);
      56            5 :             if (handleIt != it->second.channelHandles.end()) {
      57            5 :                 it->second.channelHandles.erase(handleIt);
      58            5 :                 if (it->second.channelCount > 0) {
      59            5 :                     it->second.channelCount--;
      60              :                 }
      61            5 :                 HCCL_INFO(
      62              :                     "[%s] unregistered channel[0x%llx] from endpointHandle[%p], remaining[%u].", __func__, channels[i],
      63              :                     it->first, it->second.channelCount);
      64              :                 // 注：共享 jetty 引用计数由 connection 析构时的 releaseCb_ 自动减（Endpoint::ReleaseSharedJetty），
      65              :                 // 此处不再重复减引用，仅维护 channelHandles 记录供 CheckEndpointDestroy 校验
      66            5 :                 if (it->second.channelCount == 0) {
      67            2 :                     EndpointHandle epHandle = it->first;
      68            2 :                     contexts_.erase(it);
      69            2 :                     HCCL_INFO(
      70              :                         "[%s] all channels unregistered, context removed for endpointHandle[%p].", __func__, epHandle);
      71              :                 }
      72            5 :                 break;
      73              :             }
      74              :         }
      75              :         // 非共享 jetty 的 channel 也会走 HcommChannelDestroy，此处静默忽略，避免日志刷屏
      76              :     }
      77           24 :     return HCCL_SUCCESS;
      78           24 : }
      79              : 
      80           52 : HcclResult SharedJettyMgr::CheckEndpointDestroy(EndpointHandle endpointHandle)
      81              : {
      82           52 :     std::lock_guard<std::mutex> lock(mtx_);
      83           52 :     auto it = contexts_.find(endpointHandle);
      84           52 :     if (it == contexts_.end()) {
      85           51 :         return HCCL_SUCCESS;
      86              :     }
      87            1 :     HCCL_ERROR(
      88              :         "[%s] cannot destroy endpointHandle[%p], still has [%u] shared jetty channels.", __func__, endpointHandle,
      89              :         it->second.channelCount);
      90            1 :     return HCCL_E_UNAVAIL;
      91           52 : }
      92              : 
      93            7 : bool SharedJettyMgr::HasContext(EndpointHandle endpointHandle)
      94              : {
      95            7 :     std::lock_guard<std::mutex> lock(mtx_);
      96           14 :     return contexts_.find(endpointHandle) != contexts_.end();
      97            7 : }
      98              : 
      99              : } // namespace hcomm