-
Notifications
You must be signed in to change notification settings - Fork 149
/
token.go
82 lines (68 loc) · 1.72 KB
/
token.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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package client
import (
"crypto/md5"
"encoding/hex"
"encoding/json"
"os"
"strconv"
"sync"
"github.com/google/uuid"
"k8s.io/utils/lru"
)
type IdempotentKeyGen interface {
GenerateKey(paramHash string) string
PutBack(paramHash string, uuid string)
}
// SimpleIdempotentKeyGenerator implements the generation and management of idempotency keys.
type SimpleIdempotentKeyGenerator struct {
mu sync.Mutex
cache *lru.Cache
}
func NewIdempotentKeyGenerator() *SimpleIdempotentKeyGenerator {
size, err := strconv.Atoi(os.Getenv("IDEMPOTENT_KEY_CACHE_SIZE"))
if err != nil || size <= 0 {
size = 500
}
return &SimpleIdempotentKeyGenerator{
cache: lru.New(size),
}
}
// GenerateKey generates an idempotency key based on the given parameter hash.
// multiple key is supported
func (g *SimpleIdempotentKeyGenerator) GenerateKey(paramHash string) string {
g.mu.Lock()
defer g.mu.Unlock()
v, ok := g.cache.Get(paramHash)
if ok {
uuids := v.([]string)
if len(uuids) > 0 {
id := ""
id, uuids = uuids[len(uuids)-1], uuids[:len(uuids)-1]
if len(uuids) == 0 {
g.cache.Remove(paramHash)
} else {
g.cache.Add(paramHash, uuids)
}
return id
}
}
return uuid.NewString()
}
// PutBack adds the specified idempotency key back into the cache for reuse, associating it with the given parameter hash.
func (g *SimpleIdempotentKeyGenerator) PutBack(paramHash string, uuid string) {
g.mu.Lock()
defer g.mu.Unlock()
v, ok := g.cache.Get(paramHash)
if ok {
uuids := v.([]string)
uuids = append(uuids, uuid)
g.cache.Add(paramHash, uuids)
} else {
g.cache.Add(paramHash, []string{uuid})
}
}
func md5Hash(obj any) string {
out, _ := json.Marshal(obj)
hash := md5.Sum(out)
return hex.EncodeToString(hash[:])
}