github.com/polarismesh/polaris@v1.17.8/plugin/ratelimit/lrurate/base.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 lrurate
    19  
    20  import (
    21  	"hash/crc32"
    22  
    23  	lru "github.com/hashicorp/golang-lru"
    24  	"golang.org/x/time/rate"
    25  )
    26  
    27  var (
    28  	ipLruCache      *lru.Cache
    29  	serviceLruCache *lru.Cache
    30  )
    31  
    32  // initEnv 初始化lru组件
    33  func initEnv() error {
    34  	var err error
    35  
    36  	ipLruCache, err = lru.New(rateLimitIPLruSize)
    37  	if err != nil {
    38  		return err
    39  	}
    40  
    41  	serviceLruCache, err = lru.New(rateLimitServiceLruSize)
    42  	if err != nil {
    43  		return err
    44  	}
    45  
    46  	return nil
    47  }
    48  
    49  // hash crc32取字符串hash值
    50  func hash(str string) uint32 {
    51  	return crc32.ChecksumIEEE([]byte(str))
    52  }
    53  
    54  // allowIP ip限流
    55  func allowIP(id string) bool {
    56  	key := hash(id)
    57  	ipLruCache.ContainsOrAdd(key, rate.NewLimiter(rate.Limit(rateLimitIPRate), rateLimitIPBurst))
    58  	if value, ok := ipLruCache.Get(key); ok {
    59  		return value.(*rate.Limiter).Allow()
    60  	}
    61  
    62  	return true
    63  }
    64  
    65  // allowService service限流
    66  func allowService(id string) bool {
    67  	key := hash(id)
    68  	serviceLruCache.ContainsOrAdd(key, rate.NewLimiter(rate.Limit(rateLimitServiceRate), rateLimitServiceBurst))
    69  	if value, ok := serviceLruCache.Get(key); ok {
    70  		return value.(*rate.Limiter).Allow()
    71  	}
    72  
    73  	return true
    74  }