package github
import "time"
// Token struct
type Token struct {
Hash string
RetryAfter int64
ExceededTime time.Time
}
// Tokens is the internal struct to manage the current token
// and the pool
type Tokens struct {
current int
pool []Token
}
// NewTokenManager initialize the tokens pool
func NewTokenManager(keys []string) *Tokens {
pool := []Token{}
for _, key := range keys {
t := Token{Hash: key, ExceededTime: time.Time{}, RetryAfter: 0}
pool = append(pool, t)
}
return &Tokens{
current: 0,
pool: pool,
}
}
func (r *Tokens) setCurrentTokenExceeded(retryAfter int64) {
if r.current >= len(r.pool) {
r.current %= len(r.pool)
}
if r.pool[r.current].RetryAfter == 0 {
r.pool[r.current].ExceededTime = time.Now()
r.pool[r.current].RetryAfter = retryAfter
}
}
// Get returns a new token from the token pool
func (r *Tokens) Get() *Token {
resetExceededTokens(r)
if r.current >= len(r.pool) {
r.current %= len(r.pool)
}
result := &r.pool[r.current]
r.current++
return result
}
func resetExceededTokens(r *Tokens) {
for i, token := range r.pool {
if token.RetryAfter > 0 {
if int64(time.Since(token.ExceededTime)/time.Second) > token.RetryAfter {
r.pool[i].ExceededTime = time.Time{}
r.pool[i].RetryAfter = 0
}
}
}
}