-
Notifications
You must be signed in to change notification settings - Fork 0
/
store_local.go
52 lines (43 loc) · 1.03 KB
/
store_local.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
package ratelimit
import (
"context"
"reflect"
"sync"
)
type LocalStore[T record] struct {
mu sync.RWMutex
store map[string]*T
}
func NewLocalStore[T record]() *LocalStore[T] {
return &LocalStore[T]{
store: map[string]*T{},
}
}
func (ls *LocalStore[T]) Increment(ctx context.Context, key string) error {
ls.mu.Lock()
defer ls.mu.Unlock()
val := ls.store[key]
switch any(val).(type) {
case *slidingWindowRecord:
r := reflect.ValueOf(val).Interface().(*slidingWindowRecord)
r.CurrentCount += 1
case *tokenBucketRecord:
r := reflect.ValueOf(val).Interface().(*tokenBucketRecord)
r.Count += 1
case *fixedWindowRecord:
r := reflect.ValueOf(val).Interface().(*fixedWindowRecord)
r.Count += 1
}
return nil
}
func (ls *LocalStore[T]) Get(ctx context.Context, key string) (record *T, err error) {
ls.mu.RLock()
defer ls.mu.RUnlock()
return ls.store[key], nil
}
func (ls *LocalStore[T]) Set(ctx context.Context, key string, record *T) error {
ls.mu.Lock()
defer ls.mu.Unlock()
ls.store[key] = record
return nil
}