github.com/polarismesh/polaris@v1.17.8/plugin/ratelimit/token/resource_limiter.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  	"fmt"
    22  
    23  	lru "github.com/hashicorp/golang-lru"
    24  	"golang.org/x/time/rate"
    25  
    26  	"github.com/polarismesh/polaris/plugin"
    27  )
    28  
    29  // 资源限制器
    30  type resourceRatelimit struct {
    31  	typStr    string
    32  	resources *lru.Cache
    33  	whiteList map[string]bool
    34  	config    *ResourceLimitConfig
    35  }
    36  
    37  // 新建资源限制器
    38  func newResourceRatelimit(typ plugin.RatelimitType, config *ResourceLimitConfig) (*resourceRatelimit, error) {
    39  	r := &resourceRatelimit{typStr: plugin.RatelimitStr[typ]}
    40  	if err := r.initialize(config); err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	return r, nil
    45  }
    46  
    47  // initialize
    48  func (r *resourceRatelimit) initialize(config *ResourceLimitConfig) error {
    49  	r.config = config
    50  	if config == nil || !config.Open {
    51  		log.Infof("[Plugin][%s] resource(%s) ratelimit is not open", PluginName, r.typStr)
    52  		return nil
    53  	}
    54  
    55  	if config.Global == nil {
    56  		return fmt.Errorf("resource(%s) global ratelimit rule is empty", r.typStr)
    57  	}
    58  	if config.Global.Bucket <= 0 || config.Global.Rate <= 0 {
    59  		return fmt.Errorf("resource(%s) ratelimit global bucket or rate invalid", r.typStr)
    60  	}
    61  	if config.MaxResourceCacheAmount <= 0 {
    62  		return fmt.Errorf("resource(%s) max resource amount is invalid", r.typStr)
    63  	}
    64  
    65  	cache, err := lru.New(config.MaxResourceCacheAmount)
    66  	if err != nil {
    67  		log.Errorf("[Plugin][%s] resource(%s) ratelimit create new lru cache err: %s",
    68  			PluginName, r.typStr, err.Error())
    69  		return err
    70  	}
    71  	r.resources = cache
    72  
    73  	r.whiteList = make(map[string]bool)
    74  	for _, item := range config.WhiteList {
    75  		r.whiteList[item] = true
    76  	}
    77  
    78  	log.Infof("[Plugin][%s] resource(%s) ratelimit open", PluginName, r.typStr)
    79  	return nil
    80  }
    81  
    82  // 限流是否开启
    83  func (r *resourceRatelimit) isOpen() bool {
    84  	return r.config != nil && r.config.Open
    85  }
    86  
    87  // 检查是否属于白名单，属于的话，则不限流
    88  func (r *resourceRatelimit) isWhiteList(key string) bool {
    89  	_, ok := r.whiteList[key]
    90  	return ok
    91  }
    92  
    93  // 实现limiter
    94  func (r *resourceRatelimit) allow(key string) bool {
    95  	if ok := r.isOpen(); !ok {
    96  		return true
    97  	}
    98  	if ok := r.isWhiteList(key); ok {
    99  		return true
   100  	}
   101  
   102  	value, ok := r.resources.Get(key)
   103  	if !ok {
   104  		r.resources.ContainsOrAdd(key,
   105  			rate.NewLimiter(rate.Limit(r.config.Global.Rate), r.config.Global.Bucket))
   106  		// 上面已经加了value，这里正常情况会有value
   107  		value, ok = r.resources.Get(key)
   108  		if !ok {
   109  			// 还找不到，打印日志，返回true
   110  			log.Warnf("[Plugin][%s] not found the resources(%s) key(%s) in the cache",
   111  				PluginName, r.typStr, key)
   112  			return true
   113  		}
   114  	}
   115  
   116  	return value.(*rate.Limiter).Allow()
   117  }