cubefs

1
// Copyright 2023 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package ratelimit
16

17
import (
18
	"sync"
19
	"time"
20

21
	"gopkg.in/bsm/ratelimit.v1"
22
)
23

24
type rateLimitWrapper struct {
25
	r        *ratelimit.RateLimiter
26
	refCount int
27
}
28

29
func newRateLimitWrapper(rate int) *rateLimitWrapper {
30
	return &rateLimitWrapper{
31
		r:        ratelimit.New(10*rate, 10*time.Second),
32
		refCount: 0,
33
	}
34
}
35

36
type KeyRateLimit struct {
37
	mutex   sync.RWMutex
38
	current map[string]*rateLimitWrapper // uid -> ratelimit
39
}
40

41
func NewKeyRateLimit() *KeyRateLimit {
42
	return &KeyRateLimit{current: make(map[string]*rateLimitWrapper)}
43
}
44

45
func (k *KeyRateLimit) Acquire(key string, rate int) *ratelimit.RateLimiter {
46
	k.mutex.Lock()
47
	limit, ok := k.current[key]
48
	if !ok {
49
		limit = newRateLimitWrapper(rate)
50
		k.current[key] = limit
51
	}
52
	limit.refCount++
53
	k.mutex.Unlock()
54
	return limit.r
55
}
56

57
func (k *KeyRateLimit) Release(key string) {
58
	k.mutex.Lock()
59
	defer k.mutex.Unlock()
60
	limit, ok := k.current[key]
61
	if !ok {
62
		panic("key not in map. Possible reason: Release without Acquire.")
63
	}
64
	limit.refCount--
65
	if limit.refCount < 0 {
66
		panic("internal error: refs < 0")
67
	}
68
	if limit.refCount == 0 {
69
		delete(k.current, key)
70
	}
71
}
72