feat: fetch torznab results cocurrently

This commit is contained in:
Simon Ding
2024-07-24 15:55:33 +08:00
parent 33f0a5b53f
commit a7dfa2d0f0
2 changed files with 37 additions and 14 deletions

View File

@@ -1,12 +1,14 @@
package torznab package torznab
import ( import (
"context"
"encoding/xml" "encoding/xml"
"io" "io"
"net/http" "net/http"
"net/url" "net/url"
"polaris/log" "polaris/log"
"strconv" "strconv"
"time"
"github.com/pkg/errors" "github.com/pkg/errors"
) )
@@ -59,7 +61,6 @@ type Item struct {
Name string `xml:"name,attr"` Name string `xml:"name,attr"`
Value string `xml:"value,attr"` Value string `xml:"value,attr"`
} `xml:"attr"` } `xml:"attr"`
} }
func (i *Item) GetAttr(key string) string { func (i *Item) GetAttr(key string) string {
@@ -75,12 +76,12 @@ func (r *Response) ToResults() []Result {
for _, item := range r.Channel.Item { for _, item := range r.Channel.Item {
r := Result{ r := Result{
Name: item.Title, Name: item.Title,
Magnet: item.Link, Magnet: item.Link,
Size: mustAtoI(item.Size), Size: mustAtoI(item.Size),
Seeders: mustAtoI(item.GetAttr("seeders")), Seeders: mustAtoI(item.GetAttr("seeders")),
Peers: mustAtoI(item.GetAttr("peers")), Peers: mustAtoI(item.GetAttr("peers")),
Category: mustAtoI(item.GetAttr("category")), Category: mustAtoI(item.GetAttr("category")),
Source: r.Channel.Title, Source: r.Channel.Title,
} }
res = append(res, r) res = append(res, r)
} }
@@ -96,7 +97,10 @@ func mustAtoI(key string) int {
return i return i
} }
func Search(torznabUrl, api, keyWord string) ([]Result, error) { func Search(torznabUrl, api, keyWord string) ([]Result, error) {
req, err := http.NewRequest(http.MethodGet, torznabUrl, nil) ctx, cancel := context.WithTimeout(context.TODO(), 10*time.Second)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, torznabUrl, nil)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "new request") return nil, errors.Wrap(err, "new request")
} }
@@ -130,5 +134,5 @@ type Result struct {
Seeders int Seeders int
Peers int Peers int
Category int Category int
Source string Source string
} }

View File

@@ -10,6 +10,7 @@ import (
"sort" "sort"
"strconv" "strconv"
"strings" "strings"
"sync"
"github.com/pkg/errors" "github.com/pkg/errors"
) )
@@ -97,14 +98,32 @@ func searchWithTorznab(db *db.Client, q string) []torznab.Result {
var res []torznab.Result var res []torznab.Result
allTorznab := db.GetAllTorznabInfo() allTorznab := db.GetAllTorznabInfo()
resChan := make(chan []torznab.Result)
var wg sync.WaitGroup
for _, tor := range allTorznab { for _, tor := range allTorznab {
resp, err := torznab.Search(tor.URL, tor.ApiKey, q) wg.Add(1)
if err != nil { go func () {
log.Errorf("search %s error: %v", tor.Name, err) log.Debugf("search torznab %v with %v", tor.Name, q)
continue defer wg.Done()
} resp, err := torznab.Search(tor.URL, tor.ApiKey, q)
res = append(res, resp...) if err != nil {
log.Errorf("search %s error: %v", tor.Name, err)
return
}
resChan <- resp
}()
} }
go func() {
wg.Wait()
close(resChan) // 在所有的worker完成后关闭Channel
}()
for result := range resChan {
res = append(res, result...)
}
sort.Slice(res, func(i, j int) bool { sort.Slice(res, func(i, j int) bool {
var s1 = res[i] var s1 = res[i]
var s2 = res[j] var s2 = res[j]