github.com/polarismesh/polaris@v1.17.8/common/conn/limit/api.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package connlimit
    19  
    20  import (
    21  	"errors"
    22  	"sync"
    23  
    24  	"github.com/polarismesh/polaris/common/log"
    25  )
    26  
    27  // limitListener limit obj for Listener
    28  type limitListener struct {
    29  	listenerMap map[string]*Listener // 对象索引
    30  	mu          sync.RWMutex         // 对象锁
    31  }
    32  
    33  var (
    34  	limitEntry = limitListener{
    35  		listenerMap: make(map[string]*Listener),
    36  	}
    37  )
    38  
    39  // GetLimitListener 获取当前的listener
    40  func GetLimitListener(protocol string) *Listener {
    41  	limitEntry.mu.RLock()
    42  	defer limitEntry.mu.RUnlock()
    43  	obj, ok := limitEntry.listenerMap[protocol]
    44  	if !ok {
    45  		return nil
    46  	}
    47  
    48  	return obj
    49  }
    50  
    51  // SetLimitListener 设置当前的listener
    52  // 注意：Listener.protocol不能重复
    53  func SetLimitListener(lis *Listener) error {
    54  	limitEntry.mu.Lock()
    55  	defer limitEntry.mu.Unlock()
    56  
    57  	if _, ok := limitEntry.listenerMap[lis.protocol]; ok {
    58  		log.Errorf("[ConnLimit] protocol(%s) is existed", lis.protocol)
    59  		return errors.New("protocol is existed")
    60  	}
    61  
    62  	limitEntry.listenerMap[lis.protocol] = lis
    63  	return nil
    64  }
    65  
    66  // RemoveLimitListener 清理对应协议的链接计数
    67  func RemoveLimitListener(protocol string) {
    68  	limitEntry.mu.Lock()
    69  	defer limitEntry.mu.Unlock()
    70  
    71  	delete(limitEntry.listenerMap, protocol)
    72  }