Files
zurg/pkg/realdebrid/token_manager.go
Ben Adrian Sarmiento 962845fb81 Multi-token support
2024-06-28 04:47:43 +02:00

65 lines
1.4 KiB
Go

package realdebrid
import (
"fmt"
"sync"
)
type Token struct {
value string
expired bool
}
type DownloadTokenManager struct {
tokens []Token
current int
mu sync.Mutex
}
// NewDownloadTokenManager initializes a new DownloadTokenManager with the given tokens.
func NewDownloadTokenManager(tokenStrings []string) *DownloadTokenManager {
tokens := make([]Token, len(tokenStrings))
for i, t := range tokenStrings {
tokens[i] = Token{value: t, expired: false}
}
return &DownloadTokenManager{tokens: tokens, current: 0}
}
// GetCurrentToken returns the current non-expired token.
func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
dtm.mu.Lock()
defer dtm.mu.Unlock()
for {
if !dtm.tokens[dtm.current].expired {
return dtm.tokens[dtm.current].value, nil
}
dtm.current = (dtm.current + 1) % len(dtm.tokens)
if dtm.current == 0 {
return "", fmt.Errorf("all tokens are bandwidth-limited")
}
}
}
// SetCurrentTokenExpired sets the current token as expired.
func (dtm *DownloadTokenManager) SetCurrentTokenExpired() {
dtm.mu.Lock()
defer dtm.mu.Unlock()
dtm.tokens[dtm.current].expired = true
dtm.current = (dtm.current + 1) % len(dtm.tokens)
}
// ResetAllTokens resets all tokens to expired=false.
func (dtm *DownloadTokenManager) ResetAllTokens() {
dtm.mu.Lock()
defer dtm.mu.Unlock()
for i := range dtm.tokens {
dtm.tokens[i].expired = false
}
dtm.current = 0
}