hpz937/zurg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Single worker pool, client adjustments

Browse Source
This commit is contained in:
Ben Sarmiento
2023-11-30 01:03:31 +01:00
parent 9e3760f275
commit 8de52786ce
3 changed files with 74 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
internal/config/types.go
View File

@@ -20,8 +20,6 @@ type ConfigInterface interface {
GetRandomPreferredHost() string
ShouldServeFromRclone() bool
ShouldForceIPv6() bool
GetUnrestrictWorkers() int
GetReleaseUnrestrictAfterMs() int
GetRateLimitSleepSeconds() int
GetRealDebridTimeout() int
GetRetriesUntilFailed() int
@@ -33,8 +31,6 @@ type ZurgConfig struct {
Host string `yaml:"host"`
Port string `yaml:"port"`
NumOfWorkers int `yaml:"concurrent_workers"`
UnrestrictWorkers int `yaml:"unrestrict_workers"`
ReleaseUnrestrictAfterMs int `yaml:"release_unrestrict_after_ms"`
RateLimitSleepSeconds int `yaml:"rate_limit_sleep_secs"`
RefreshEverySeconds int `yaml:"check_for_changes_every_secs"`
CanRepair bool `yaml:"enable_repair"`
@@ -128,20 +124,6 @@ func (z *ZurgConfig) ShouldForceIPv6() bool {
return z.ForceIPv6
}
func (z *ZurgConfig) GetUnrestrictWorkers() int {
if z.UnrestrictWorkers == 0 {
return 20
}
return z.UnrestrictWorkers
}
func (z *ZurgConfig) GetReleaseUnrestrictAfterMs() int {
if z.ReleaseUnrestrictAfterMs == 0 {
return 100
}
return z.ReleaseUnrestrictAfterMs
}
func (z *ZurgConfig) GetRateLimitSleepSeconds() int {
if z.RateLimitSleepSeconds == 0 {
return 4

16
internal/torrent/manager.go
View File

@@ -36,7 +36,6 @@ type TorrentManager struct {
latestState *LibraryState
requiredVersion string
workerPool *ants.Pool
unrestrictPool *ants.Pool
log *zap.SugaredLogger
}
@@ -57,14 +56,6 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
log: log,
}
// create unrestrict pool
unrestrictPool, err := ants.NewPool(t.Config.GetUnrestrictWorkers())
if err != nil {
t.unrestrictPool = t.workerPool
} else {
t.unrestrictPool = unrestrictPool
}
// create internal directories
t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is AccessKey
t.DirectoryMap.Set(INT_INFO_CACHE, cmap.New[*Torrent]()) // key is Torrent ID
@@ -92,6 +83,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
})
var newTorrents []realdebrid.Torrent
var err error
initWait.Add(1)
_ = t.workerPool.Submit(func() {
defer initWait.Done()
@@ -212,9 +204,8 @@ func (t *TorrentManager) mergeToMain(mainTorrent, torrentToMerge *Torrent) *Torr
// proxy
func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download {
retChan := make(chan *realdebrid.Download, 1)
t.unrestrictPool.Submit(func() {
t.workerPool.Submit(func() {
retChan <- t.Api.UnrestrictUntilOk(link, t.Config.ShouldServeFromRclone())
time.Sleep(time.Duration(t.Config.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
})
defer close(retChan)
return <-retChan
@@ -551,7 +542,7 @@ func (t *TorrentManager) organizeChaos(links []string, selectedFiles []*File) ([
wg.Add(1)
link := link // redeclare to avoid closure on loop variable
// Use the existing worker pool to submit tasks
_ = t.unrestrictPool.Submit(func() {
_ = t.workerPool.Submit(func() {
defer wg.Done()
if t.DownloadCache.Has(link) {
download, _ := t.DownloadCache.Get(link)
@@ -560,7 +551,6 @@ func (t *TorrentManager) organizeChaos(links []string, selectedFiles []*File) ([
}
resp := t.Api.UnrestrictUntilOk(link, t.Config.ShouldServeFromRclone())
resultsChan <- Result{Response: resp}
time.Sleep(time.Duration(t.Config.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
})
}

132
pkg/http/client.go
View File

@@ -13,38 +13,79 @@ import (
"go.uber.org/zap"
)
const (
RATE_LIMIT_FACTOR = 4 // should always be > 1
)
type HTTPClient struct {
Client *http.Client
MaxRetries int
Backoff func(attempt int) time.Duration
ShouldRetry func(resp *http.Response, err error) int
BearerToken string
client *http.Client
maxRetries int
backoff func(attempt int) time.Duration
getRetryIncr func(resp *http.Response, err error) int
bearerToken string
cfg config.ConfigInterface
ipv6 cmap.ConcurrentMap[string, net.IP]
log *zap.SugaredLogger
config config.ConfigInterface
IPv6 cmap.ConcurrentMap[string, net.IP]
}
func NewHTTPClient(token string, maxRetries int, timeoutSecs int, cfg config.ConfigInterface, log *zap.SugaredLogger) *HTTPClient {
return &HTTPClient{
bearerToken: token,
client: &http.Client{
Timeout: time.Duration(timeoutSecs) * time.Second,
},
maxRetries: maxRetries * RATE_LIMIT_FACTOR,
backoff: func(attempt int) time.Duration {
maxDuration := 60
backoff := int(math.Pow(2, float64(attempt)))
if backoff > maxDuration {
backoff = maxDuration
}
return time.Duration(backoff) * time.Second
},
getRetryIncr: func(resp *http.Response, err error) int {
if resp != nil {
if resp.StatusCode == 429 || resp.StatusCode == 400 || resp.StatusCode == 403 {
return 1 // retry but don't increment attempt
}
return 0 // don't retry
} else if err != nil {
errStr := err.Error()
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "no such host") {
return 1 // retry but don't increment attempt
} else {
return RATE_LIMIT_FACTOR
}
}
return RATE_LIMIT_FACTOR // retry and increment attempt
},
cfg: cfg,
ipv6: cmap.New[net.IP](),
log: log,
}
}
func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
if r.config != nil && strings.Contains(req.Host, "download.real-debrid.") {
prefHost := r.config.GetRandomPreferredHost()
if r.cfg != nil && strings.Contains(req.Host, "download.real-debrid.") {
prefHost := r.cfg.GetRandomPreferredHost()
if prefHost != "" {
req.Host = prefHost
req.URL.Host = prefHost
}
}
if r.BearerToken != "" {
req.Header.Set("Authorization", "Bearer "+r.BearerToken)
if r.bearerToken != "" {
req.Header.Set("Authorization", "Bearer "+r.bearerToken)
}
if r.config.ShouldForceIPv6() {
if r.cfg.ShouldForceIPv6() {
dialer := &net.Dialer{}
transport := r.Client.Transport.(*http.Transport).Clone()
transport := r.client.Transport.(*http.Transport).Clone()
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, err
}
if ipv6, ok := r.IPv6.Get(host); !ok {
if ipv6, ok := r.ipv6.Get(host); !ok {
// Lookup IP address if not found in map
ips, err := net.LookupIP(host)
if err != nil {
@@ -53,7 +94,7 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
for _, ip := range ips {
if ip.To4() == nil { // IPv6
ipv6 = ip
r.IPv6.Set(host, ipv6)
r.ipv6.Set(host, ipv6)
break
}
}
@@ -67,66 +108,31 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
}
return dialer.DialContext(ctx, network, addr)
}
r.Client.Transport = transport
r.client.Transport = transport
}
var resp *http.Response
var err error
attempt := 0
for {
resp, err = r.Client.Do(req)
if val := r.ShouldRetry(resp, err); val == -1 {
return resp, err
} else {
if val == 0 {
time.Sleep(time.Duration(r.config.GetRateLimitSleepSeconds()) * time.Second) // extra delay
} else {
attempt += val
if attempt > r.MaxRetries {
return resp, err
resp, err = r.client.Do(req)
if incr := r.getRetryIncr(resp, err); incr > 0 {
attempt += incr
if attempt > r.maxRetries {
break
}
time.Sleep(r.Backoff(attempt))
if incr >= RATE_LIMIT_FACTOR {
time.Sleep(r.backoff(attempt))
} else {
time.Sleep(time.Duration(r.cfg.GetRateLimitSleepSeconds()) * time.Second) // extra delay
}
if resp != nil {
resp.Body.Close()
}
}
}
}
func NewHTTPClient(token string, maxRetries int, timeoutSecs int, cfg config.ConfigInterface, log *zap.SugaredLogger) *HTTPClient {
return &HTTPClient{
BearerToken: token,
Client: &http.Client{
Timeout: time.Duration(timeoutSecs) * time.Second,
},
MaxRetries: maxRetries,
Backoff: func(attempt int) time.Duration {
maxDuration := 60
backoff := int(math.Pow(2, float64(attempt)))
if backoff > maxDuration {
backoff = maxDuration
}
return time.Duration(backoff) * time.Second
},
ShouldRetry: func(resp *http.Response, err error) int {
if resp != nil {
if resp.StatusCode == 429 || resp.StatusCode == 400 || resp.StatusCode == 403 {
return 0 // retry but don't increment attempt
}
return -1 // don't retry
} else if err != nil {
errStr := err.Error()
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "no such host") {
return 0 // retry but don't increment attempt
} else {
return 1
// if incr == 0, don't retry anymore
break
}
}
return 1 // retry and increment attempt
},
log: log,
config: cfg,
IPv6: cmap.New[net.IP](),
}
return resp, err
}