forked from momokatte/go-limiter
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtoken.go
59 lines (52 loc) · 1.75 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
package limiter
import (
"sync"
)
/*
TokenChanLimiter enforces a concurrency limit using tokens and satisfies the TokenLimiter and InvocationLimiter interfaces.
*/
type TokenChanLimiter struct {
mu sync.Mutex
tokens chan *[16]byte
}
/*
NewTokenChanLimiter instantiates a new TokenChanLimiter with the provided number of tokens.
*/
func NewTokenChanLimiter(initialTokens uint) (l *TokenChanLimiter) {
l = &TokenChanLimiter{
tokens: make(chan *[16]byte, initialTokens),
}
fillTokenChan(l.tokens)
return
}
/*
AcquireToken blocks until a token can be acquired from the limiter's supply. The token must be held for the duration of the activity which needs to be limited, and then it must be passed to the ReleaseToken method without modification.
*/
func (l *TokenChanLimiter) AcquireToken() (token *[16]byte) {
select {
case token = <-l.tokens:
return
}
}
/*
ReleaseToken notifies the limiter that the provided token (pointer and value) can be used by another goroutine. The caller must not modify the value of the token at any time, but if the token implementation is known by the caller then unmarshaling of its value is not discouraged.
*/
func (l *TokenChanLimiter) ReleaseToken(token *[16]byte) {
l.tokens <- token
}
/*
Invoke enforces the limiter's limits around the invocation of the passed function. The error returned by the function invocation is returned to the caller without modification, and its existence may be used by the limiter to delay the current return or subsequent invocations.
*/
func (l *TokenChanLimiter) Invoke(f func() error) (err error) {
token := l.AcquireToken()
err = f()
l.ReleaseToken(token)
return
}
func fillTokenChan(c chan *[16]byte) {
capacity := cap(c)
for i := 0; i < capacity; i += 1 {
c <- new([16]byte)
}
return
}