Improve timeout logic.

This commit is contained in:
Syfaro 2019-09-19 11:25:53 -05:00
parent 8318bfa8d9
commit 5e41802ce4
2 changed files with 25 additions and 5 deletions

2
go.mod
View File

@ -42,3 +42,5 @@ require (
huefox.com/syfaro/go-e621 v1.0.0 huefox.com/syfaro/go-e621 v1.0.0
mvdan.cc/xurls v1.1.1-0.20180901190342-70405f5eab51 mvdan.cc/xurls v1.1.1-0.20180901190342-70405f5eab51
) )
go 1.13

View File

@ -8,10 +8,11 @@ import (
"os/exec" "os/exec"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/PuerkitoBio/goquery" "github.com/PuerkitoBio/goquery"
"github.com/Syfaro/finch" "github.com/Syfaro/finch"
"github.com/go-telegram-bot-api/telegram-bot-api/v5" tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
"huefox.com/syfaro/telegram-furryimgbot/logger" "huefox.com/syfaro/telegram-furryimgbot/logger"
) )
@ -27,6 +28,7 @@ type imageURLResult struct {
type fa struct { type fa struct {
cookies []*http.Cookie cookies []*http.Cookie
fautilKey string fautilKey string
client *http.Client
} }
func (fa) Name() string { func (fa) Name() string {
@ -56,12 +58,24 @@ type cfscrape struct {
UserAgent string `json:"user_agent"` UserAgent string `json:"user_agent"`
} }
func (f *fa) initClient() {
if f.client != nil {
return
}
timeout := time.Duration(4 * time.Second)
client := &http.Client{
Timeout: timeout,
}
f.client = client
}
func (f fa) handleDirectLink(postURL string) ([]PostInfo, error) { func (f fa) handleDirectLink(postURL string) ([]PostInfo, error) {
logger.Log.Debug("Attempting to load FA direct link") logger.Log.Debug("Attempting to load FA direct link")
post := strings.TrimPrefix(postURL, "https://") post := strings.TrimPrefix(postURL, "https://")
post = strings.TrimPrefix(post, "http://") post = strings.TrimPrefix(post, "http://")
post = "%" + post post = "https://" + post
logger.Log.Debugf("Converted post link is %s", post) logger.Log.Debugf("Converted post link is %s", post)
@ -74,16 +88,20 @@ func (f fa) handleDirectLink(postURL string) ([]PostInfo, error) {
req, _ := http.NewRequest("GET", u.String(), nil) req, _ := http.NewRequest("GET", u.String(), nil)
req.Header.Set("X-Api-Key", f.fautilKey) req.Header.Set("X-Api-Key", f.fautilKey)
resp, err := http.DefaultClient.Do(req) f.initClient()
resp, err := f.client.Do(req)
if err != nil { if err != nil {
return nil, nil logger.Log.Warnf("Error loading, falling back to direct loading")
d := direct{}
return d.GetImageURLs(postURL, tgbotapi.User{})
} }
defer resp.Body.Close() defer resp.Body.Close()
var results []imageURLResult var results []imageURLResult
d := json.NewDecoder(resp.Body) d := json.NewDecoder(resp.Body)
if err = d.Decode(&results); err != nil { if err = d.Decode(&results); err != nil {
logger.Log.Debug("Unable to decode results") logger.Log.Warnf("Unable to decode results")
return nil, err return nil, err
} }