github.com/polarismesh/polaris@v1.17.8/plugin/ratelimit/token/implement.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 token
    19  
    20  import (
    21  	"github.com/polarismesh/polaris/plugin"
    22  )
    23  
    24  // initialize 插件初始化函数
    25  func (tb *tokenBucket) initialize(c *plugin.ConfigEntry) error {
    26  	config, err := decodeConfig(c.Option)
    27  	if err != nil {
    28  		log.Errorf("[Plugin][%s] initialize err: %s", PluginName, err.Error())
    29  		return err
    30  	}
    31  
    32  	tb.config = config
    33  	tb.limiters = make(map[plugin.RatelimitType]limiter)
    34  
    35  	// IP限流
    36  	irt, err := newResourceRatelimit(plugin.IPRatelimit, config.IPLimitConf)
    37  	if err != nil {
    38  		return err
    39  	}
    40  	tb.limiters[plugin.IPRatelimit] = irt
    41  
    42  	// 接口限流
    43  	art, err := newAPIRatelimit(config.APILimitConf)
    44  	if err != nil {
    45  		return err
    46  	}
    47  	tb.limiters[plugin.APIRatelimit] = art
    48  
    49  	// 操作实例限流
    50  	instance, err := newResourceRatelimit(plugin.InstanceRatelimit, config.InstanceLimitConf)
    51  	if err != nil {
    52  		return err
    53  	}
    54  	tb.limiters[plugin.InstanceRatelimit] = instance
    55  
    56  	return nil
    57  }
    58  
    59  // allow 插件的限流实现函数
    60  func (tb *tokenBucket) allow(typ plugin.RatelimitType, key string) bool {
    61  	// key为空，则不作限制
    62  	if key == "" {
    63  		return true
    64  	}
    65  	l, ok := tb.limiters[typ]
    66  	if !ok {
    67  		return true
    68  	}
    69  
    70  	return l.allow(key)
    71  }