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          209 : SharedJettyMgr& SharedJettyMgr::GetInstance()
      17              : {
      18          209 :     static SharedJettyMgr instance;
      19          209 :     return instance;
      20              : }
      21              : 
      22              : HcclResult
      23           16 : SharedJettyMgr::RegisterChannels(EndpointHandle endpointHandle, const ChannelHandle* channels, uint32_t channelNum)
      24              : {
      25           16 :     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           13 :     std::lock_guard<std::mutex> lock(mtx_);
      33           13 :     auto& ctx = contexts_[endpointHandle];
      34           32 :     for (uint32_t i = 0; i < channelNum; ++i) {
      35           19 :         ctx.channelHandles.insert(channels[i]);
      36              :     }
      37              :     // 以集合实际大小为准：insert 对重复句柄为 no-op，若用 += channelNum 会导致
      38              :     // channelCount > channelHandles.size()，UnregisterChannels 永远无法将 count 减到 0，
      39              :     // CheckEndpointDestroy 永久阻塞 endpoint 销毁。
      40           13 :     ctx.channelCount = static_cast<uint32_t>(ctx.channelHandles.size());
      41           13 :     HCCL_INFO(
      42              :         "[%s] registered %u channels for endpointHandle[%p], total channelCount[%u].", __func__, channelNum,
      43              :         endpointHandle, ctx.channelCount);
      44           13 :     return HCCL_SUCCESS;
      45           13 : }
      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            5 :                 if (it->second.channelCount == 0) {
      66            2 :                     EndpointHandle epHandle = it->first;
      67            2 :                     contexts_.erase(it);
      68            2 :                     HCCL_INFO(
      69              :                         "[%s] all channels unregistered, context removed for endpointHandle[%p].", __func__, epHandle);
      70              :                 }
      71            5 :                 break;
      72              :             }
      73              :         }
      74              :         // 非共享 jetty 的 channel 也会走 HcommChannelDestroy，此处静默忽略，避免日志刷屏
      75              :     }
      76           24 :     return HCCL_SUCCESS;
      77           24 : }
      78              : 
      79           54 : HcclResult SharedJettyMgr::CheckEndpointDestroy(EndpointHandle endpointHandle)
      80              : {
      81           54 :     std::lock_guard<std::mutex> lock(mtx_);
      82           54 :     auto it = contexts_.find(endpointHandle);
      83           54 :     if (it == contexts_.end()) {
      84           53 :         return HCCL_SUCCESS;
      85              :     }
      86            1 :     HCCL_ERROR(
      87              :         "[%s] cannot destroy endpointHandle[%p], still has [%u] shared jetty channels.", __func__, endpointHandle,
      88              :         it->second.channelCount);
      89            1 :     return HCCL_E_UNAVAIL;
      90           54 : }
      91              : 
      92           12 : bool SharedJettyMgr::HasContext(EndpointHandle endpointHandle)
      93              : {
      94           12 :     std::lock_guard<std::mutex> lock(mtx_);
      95           24 :     return contexts_.find(endpointHandle) != contexts_.end();
      96           12 : }
      97              : 
      98           58 : void SharedJettyMgr::UnregisterEndpoint(EndpointHandle endpointHandle)
      99              : {
     100           58 :     std::lock_guard<std::mutex> lock(mtx_);
     101           58 :     auto it = contexts_.find(endpointHandle);
     102           58 :     if (it == contexts_.end()) {
     103           56 :         return;
     104              :     }
     105            2 :     if (it->second.channelCount > 0) {
     106            2 :         HCCL_WARNING(
     107              :             "[%s] endpointHandle[%p] still has [%u] channels, force remove to avoid handle reuse misjudge.", __func__,
     108              :             endpointHandle, it->second.channelCount);
     109              :     }
     110            2 :     contexts_.erase(it);
     111            2 :     HCCL_INFO("[%s] unregistered endpointHandle[%p] from shared jetty mgr.", __func__, endpointHandle);
     112           58 : }
     113              : 
     114              : } // namespace hcomm