Rewrote account switching, fix some minor bugs here and there
This commit is contained in:
Mukhtar Akere
71
pkg/debrid/account/account.go
Normal file
71
pkg/debrid/account/account.go
Normal file
@@ -0,0 +1,71 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/puzpuzpuz/xsync/v4"
|
||||
"github.com/sirrobot01/decypharr/internal/request"
|
||||
"github.com/sirrobot01/decypharr/pkg/debrid/types"
|
||||
)
|
||||
|
||||
type Account struct {
|
||||
Debrid string `json:"debrid"` // The debrid service name, e.g. "realdebrid"
|
||||
links *xsync.Map[string, types.DownloadLink] // key is the sliced file link
|
||||
Index int `json:"index"` // The index of the account in the config
|
||||
Disabled atomic.Bool `json:"disabled"`
|
||||
Token string `json:"token"`
|
||||
TrafficUsed atomic.Int64 `json:"traffic_used"` // Traffic used in bytes
|
||||
Username string `json:"username"` // Username for the account
|
||||
httpClient *request.Client
|
||||
}
|
||||
|
||||
func (a *Account) Equals(other *Account) bool {
|
||||
if other == nil {
|
||||
return false
|
||||
}
|
||||
return a.Token == other.Token && a.Debrid == other.Debrid
|
||||
}
|
||||
|
||||
func (a *Account) Client() *request.Client {
|
||||
return a.httpClient
|
||||
}
|
||||
|
||||
// slice download link
|
||||
func (a *Account) sliceFileLink(fileLink string) string {
|
||||
if a.Debrid != "realdebrid" {
|
||||
return fileLink
|
||||
}
|
||||
if len(fileLink) < 39 {
|
||||
return fileLink
|
||||
}
|
||||
return fileLink[0:39]
|
||||
}
|
||||
|
||||
func (a *Account) GetDownloadLink(fileLink string) (types.DownloadLink, error) {
|
||||
slicedLink := a.sliceFileLink(fileLink)
|
||||
dl, ok := a.links.Load(slicedLink)
|
||||
if !ok {
|
||||
return types.DownloadLink{}, types.ErrDownloadLinkNotFound
|
||||
}
|
||||
return dl, nil
|
||||
}
|
||||
|
||||
func (a *Account) StoreDownloadLink(dl types.DownloadLink) {
|
||||
slicedLink := a.sliceFileLink(dl.Link)
|
||||
a.links.Store(slicedLink, dl)
|
||||
}
|
||||
func (a *Account) DeleteDownloadLink(fileLink string) {
|
||||
slicedLink := a.sliceFileLink(fileLink)
|
||||
a.links.Delete(slicedLink)
|
||||
}
|
||||
func (a *Account) ClearDownloadLinks() {
|
||||
a.links.Clear()
|
||||
}
|
||||
func (a *Account) DownloadLinksCount() int {
|
||||
return a.links.Size()
|
||||
}
|
||||
func (a *Account) StoreDownloadLinks(dls map[string]*types.DownloadLink) {
|
||||
for _, dl := range dls {
|
||||
a.StoreDownloadLink(*dl)
|
||||
}
|
||||
}
|
||||
211
pkg/debrid/account/manager.go
Normal file
211
pkg/debrid/account/manager.go
Normal file
@@ -0,0 +1,211 @@
|
||||
package account
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/puzpuzpuz/xsync/v4"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/sirrobot01/decypharr/internal/config"
|
||||
"github.com/sirrobot01/decypharr/internal/request"
|
||||
"github.com/sirrobot01/decypharr/internal/utils"
|
||||
"github.com/sirrobot01/decypharr/pkg/debrid/types"
|
||||
"go.uber.org/ratelimit"
|
||||
)
|
||||
|
||||
type Manager struct {
|
||||
debrid string
|
||||
current atomic.Pointer[Account]
|
||||
accounts *xsync.Map[string, *Account]
|
||||
}
|
||||
|
||||
func NewManager(debridConf config.Debrid, downloadRL ratelimit.Limiter, logger zerolog.Logger) *Manager {
|
||||
m := &Manager{
|
||||
debrid: debridConf.Name,
|
||||
accounts: xsync.NewMap[string, *Account](),
|
||||
}
|
||||
|
||||
var firstAccount *Account
|
||||
for idx, token := range debridConf.DownloadAPIKeys {
|
||||
if token == "" {
|
||||
continue
|
||||
}
|
||||
headers := map[string]string{
|
||||
"Authorization": fmt.Sprintf("Bearer %s", token),
|
||||
}
|
||||
account := &Account{
|
||||
Debrid: debridConf.Name,
|
||||
Token: token,
|
||||
Index: idx,
|
||||
links: xsync.NewMap[string, types.DownloadLink](),
|
||||
httpClient: request.New(
|
||||
request.WithRateLimiter(downloadRL),
|
||||
request.WithLogger(logger),
|
||||
request.WithHeaders(headers),
|
||||
request.WithMaxRetries(3),
|
||||
request.WithRetryableStatus(429, 447, 502),
|
||||
request.WithProxy(debridConf.Proxy),
|
||||
),
|
||||
}
|
||||
m.accounts.Store(token, account)
|
||||
if firstAccount == nil {
|
||||
firstAccount = account
|
||||
}
|
||||
}
|
||||
m.current.Store(firstAccount)
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *Manager) Active() []*Account {
|
||||
activeAccounts := make([]*Account, 0)
|
||||
m.accounts.Range(func(key string, acc *Account) bool {
|
||||
if !acc.Disabled.Load() {
|
||||
activeAccounts = append(activeAccounts, acc)
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
slices.SortFunc(activeAccounts, func(i, j *Account) int {
|
||||
return i.Index - j.Index
|
||||
})
|
||||
return activeAccounts
|
||||
}
|
||||
|
||||
func (m *Manager) All() []*Account {
|
||||
allAccounts := make([]*Account, 0)
|
||||
m.accounts.Range(func(key string, acc *Account) bool {
|
||||
allAccounts = append(allAccounts, acc)
|
||||
return true
|
||||
})
|
||||
|
||||
slices.SortFunc(allAccounts, func(i, j *Account) int {
|
||||
return i.Index - j.Index
|
||||
})
|
||||
return allAccounts
|
||||
}
|
||||
|
||||
func (m *Manager) Current() *Account {
|
||||
// Fast path - most common case
|
||||
current := m.current.Load()
|
||||
if current != nil && !current.Disabled.Load() {
|
||||
return current
|
||||
}
|
||||
|
||||
// Slow path - find new current account
|
||||
activeAccounts := m.Active()
|
||||
if len(activeAccounts) == 0 {
|
||||
m.current.Store(nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
newCurrent := activeAccounts[0]
|
||||
m.current.Store(newCurrent)
|
||||
return newCurrent
|
||||
}
|
||||
|
||||
func (m *Manager) setCurrent(account *Account) {
|
||||
m.current.Store(account)
|
||||
}
|
||||
|
||||
func (m *Manager) Disable(account *Account) {
|
||||
if account == nil {
|
||||
return
|
||||
}
|
||||
|
||||
account.Disabled.Store(true)
|
||||
|
||||
// If we're disabling the current account, it will be replaced
|
||||
// on the next Current() call - no need to proactively update
|
||||
current := m.current.Load()
|
||||
if current != nil && current.Token == account.Token {
|
||||
// Optional: immediately find replacement
|
||||
activeAccounts := m.Active()
|
||||
if len(activeAccounts) > 0 {
|
||||
m.current.Store(activeAccounts[0])
|
||||
} else {
|
||||
m.current.Store(nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) Reset() {
|
||||
m.accounts.Range(func(key string, acc *Account) bool {
|
||||
acc.Disabled.Store(false)
|
||||
return true
|
||||
})
|
||||
|
||||
// Set current to first active account
|
||||
activeAccounts := m.Active()
|
||||
if len(activeAccounts) > 0 {
|
||||
m.current.Store(activeAccounts[0])
|
||||
} else {
|
||||
m.current.Store(nil)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) Update(account *Account) {
|
||||
if account != nil {
|
||||
m.accounts.Store(account.Token, account)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) GetAccount(token string) (*Account, error) {
|
||||
if token == "" {
|
||||
return nil, fmt.Errorf("token cannot be empty")
|
||||
}
|
||||
acc, ok := m.accounts.Load(token)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("account not found for token")
|
||||
}
|
||||
return acc, nil
|
||||
}
|
||||
|
||||
func (m *Manager) GetDownloadLink(fileLink string) (types.DownloadLink, error) {
|
||||
current := m.Current()
|
||||
if current == nil {
|
||||
return types.DownloadLink{}, fmt.Errorf("no active account for debrid service %s", m.debrid)
|
||||
}
|
||||
return current.GetDownloadLink(fileLink)
|
||||
}
|
||||
|
||||
func (m *Manager) GetAccountFromDownloadLink(downloadLink types.DownloadLink) (*Account, error) {
|
||||
if downloadLink.Link == "" {
|
||||
return nil, fmt.Errorf("cannot get account from empty download link")
|
||||
}
|
||||
if downloadLink.Token == "" {
|
||||
return nil, fmt.Errorf("cannot get account from download link without token")
|
||||
}
|
||||
return m.GetAccount(downloadLink.Token)
|
||||
}
|
||||
|
||||
func (m *Manager) StoreDownloadLink(downloadLink types.DownloadLink) {
|
||||
if downloadLink.Link == "" || downloadLink.Token == "" {
|
||||
return
|
||||
}
|
||||
account, err := m.GetAccount(downloadLink.Token)
|
||||
if err != nil || account == nil {
|
||||
return
|
||||
}
|
||||
account.StoreDownloadLink(downloadLink)
|
||||
}
|
||||
|
||||
func (m *Manager) Stats() []map[string]any {
|
||||
stats := make([]map[string]any, 0)
|
||||
|
||||
for _, acc := range m.All() {
|
||||
maskedToken := utils.Mask(acc.Token)
|
||||
accountDetail := map[string]any{
|
||||
"in_use": acc.Equals(m.Current()),
|
||||
"order": acc.Index,
|
||||
"disabled": acc.Disabled.Load(),
|
||||
"token_masked": maskedToken,
|
||||
"username": acc.Username,
|
||||
"traffic_used": acc.TrafficUsed.Load(),
|
||||
"links_count": acc.DownloadLinksCount(),
|
||||
"debrid": acc.Debrid,
|
||||
}
|
||||
stats = append(stats, accountDetail)
|
||||
}
|
||||
return stats
|
||||
}
|
||||
Reference in New Issue
Block a user