Add job for monitoring bw limit status of tokens

This commit is contained in:
Ben Adrian Sarmiento
2024-06-28 18:55:02 +02:00
parent c3aea427d0
commit 67111696a2
13 changed files with 136 additions and 55 deletions

View File

@@ -38,7 +38,7 @@ func NewRealDebrid(apiClient, unrestrictClient, downloadClient *zurghttp.HTTPCli
rd := &RealDebrid{
UnrestrictMap: cmap.New[cmap.ConcurrentMap[string, *Download]](),
TokenManager: NewDownloadTokenManager(downloadTokens),
TokenManager: NewDownloadTokenManager(downloadTokens, log),
torrentsCache: []Torrent{},
verifiedLinks: cmap.New[int64](),
apiClient: apiClient,
@@ -69,8 +69,8 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
}
// check if the link is already unrestricted
tokenMap, ok := rd.UnrestrictMap.Get(token)
if ok && tokenMap.Has(link) {
tokenMap, _ := rd.UnrestrictMap.Get(token)
if tokenMap.Has(link) {
download, _ := tokenMap.Get(link)
// check if the link is in the verified links cache
if expiry, ok := rd.verifiedLinks.Get(download.ID); ok && expiry > time.Now().Unix() {
@@ -78,8 +78,8 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
}
err := rd.downloadClient.VerifyLink(download.Download)
if utils.IsBWLimitExceeded(err) {
rd.TokenManager.SetTokenAsExpired(token)
if utils.IsBytesLimitReached(err) {
rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
continue
}
if err != nil {
@@ -94,12 +94,13 @@ func (rd *RealDebrid) UnrestrictAndVerify(link string) (*Download, error) {
if err != nil {
return nil, err
}
download.Token = token
tokenMap.Set(link, download)
err = rd.downloadClient.VerifyLink(download.Download)
if utils.IsBWLimitExceeded(err) {
rd.TokenManager.SetTokenAsExpired(token)
if utils.IsBytesLimitReached(err) {
rd.TokenManager.SetTokenAsExpired(token, "bandwidth limit exceeded")
continue
}
if err != nil {
@@ -404,3 +405,35 @@ func (rd *RealDebrid) AvailabilityCheck(hashes []string) (AvailabilityResponse,
func (rd *RealDebrid) DownloadFile(req *http.Request) (*http.Response, error) {
return rd.downloadClient.Do(req)
}
// MonitorExpiredTokens is a permanent job for monitoring expired tokens if they are still expired
func (rd *RealDebrid) MonitorExpiredTokens() {
sleepPeriod := 5 * time.Minute
rd.workerPool.Submit(func() {
for {
expiredTokens := rd.TokenManager.GetExpiredTokens()
for _, token := range expiredTokens {
tokenMap, _ := rd.UnrestrictMap.Get(token)
stillExpired := true
skipAll := false
tokenMap.IterCb(func(key string, download *Download) {
if skipAll {
return
}
err := rd.downloadClient.VerifyLink(download.Download)
if err != nil {
skipAll = utils.IsBytesLimitReached(err)
return
}
stillExpired = false
skipAll = true
rd.verifiedLinks.Set(download.ID, time.Now().Unix()+60*60*24)
})
if !stillExpired {
rd.TokenManager.SetTokenAsUnexpired(token)
}
}
time.Sleep(sleepPeriod)
}
})
}

View File

@@ -3,6 +3,9 @@ package realdebrid
import (
"fmt"
"sync"
"github.com/debridmediamanager/zurg/pkg/logutil"
"github.com/debridmediamanager/zurg/pkg/utils"
)
type Token struct {
@@ -14,15 +17,16 @@ type DownloadTokenManager struct {
tokens []Token
current int
mu sync.RWMutex
log *logutil.Logger
}
// NewDownloadTokenManager initializes a new DownloadTokenManager with the given tokens.
func NewDownloadTokenManager(tokenStrings []string) *DownloadTokenManager {
func NewDownloadTokenManager(tokenStrings []string, log *logutil.Logger) *DownloadTokenManager {
tokens := make([]Token, len(tokenStrings))
for i, t := range tokenStrings {
tokens[i] = Token{value: t, expired: false}
}
return &DownloadTokenManager{tokens: tokens, current: 0}
return &DownloadTokenManager{tokens: tokens, current: 0, log: log}
}
// GetCurrentToken returns the current non-expired token.
@@ -38,19 +42,36 @@ func (dtm *DownloadTokenManager) GetCurrentToken() (string, error) {
dtm.current = (dtm.current + 1) % len(dtm.tokens)
if dtm.current == 0 {
return "", fmt.Errorf("all tokens are bandwidth-limited")
return "", fmt.Errorf("all tokens are expired")
}
}
}
// SetTokenAsExpired sets the specified token as expired.
func (dtm *DownloadTokenManager) SetTokenAsExpired(token string) error {
func (dtm *DownloadTokenManager) SetTokenAsExpired(token, reason string) error {
dtm.mu.Lock()
defer dtm.mu.Unlock()
for i, t := range dtm.tokens {
if t.value == token {
dtm.tokens[i].expired = true
dtm.log.Debugf("Token %s set as expired (reason=%s)", utils.MaskToken(token), reason)
return nil
}
}
return fmt.Errorf("token not found")
}
// SetTokenAsUnexpired sets the specified token as unexpired.
func (dtm *DownloadTokenManager) SetTokenAsUnexpired(token string) error {
dtm.mu.Lock()
defer dtm.mu.Unlock()
for i, t := range dtm.tokens {
if t.value == token {
dtm.tokens[i].expired = false
dtm.log.Debugf("Token %s set as unexpired", utils.MaskToken(token))
return nil
}
}
@@ -68,3 +89,16 @@ func (dtm *DownloadTokenManager) ResetAllTokens() {
}
dtm.current = 0
}
func (dtm *DownloadTokenManager) GetExpiredTokens() []string {
dtm.mu.RLock()
defer dtm.mu.RUnlock()
var tokens []string
for _, t := range dtm.tokens {
if t.expired {
tokens = append(tokens, t.value)
}
}
return tokens
}

View File

@@ -24,6 +24,7 @@ type Download struct {
Download string `json:"download"` // Generated link
Streamable int `json:"streamable"`
Generated string `json:"-"` // jsonDate
Token string `json:"-"`
}
func (d *Download) UnmarshalJSON(data []byte) error {