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