diff --git a/README.md b/README.md index 98d6152..1b3683e 100644 --- a/README.md +++ b/README.md @@ -255,7 +255,7 @@ This project is licensed under the MIT License. See [LICENSE](LICENSE). ## 🤝 Contributing -Contributions are welcome. Please feel free to submit a pull request. +Contributions are welcome. See [CONTRIBUTE](./docs/CONTRIBUTING.md). Please feel free to submit your improvements! ## 👾 Issues & Support diff --git a/baidu/search.go b/baidu/search.go index 545c0ef..e744746 100644 --- a/baidu/search.go +++ b/baidu/search.go @@ -1,6 +1,7 @@ package baidu import ( + "context" "encoding/json" "fmt" "regexp" @@ -70,7 +71,8 @@ func (baid *Baidu) isTimeout(page *rod.Page) bool { // Search executes a Baidu web search and returns normalized search results. // It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (baid *Baidu) Search(query core.Query) ([]core.SearchResult, error) { +func (baid *Baidu) Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) baid.logger.Debug("Starting search, query: %+v", query) searchResults := []core.SearchResult{} @@ -81,7 +83,7 @@ func (baid *Baidu) Search(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := baid.Navigate(url) + page, err := baid.Navigate(ctx, url) if err != nil { return nil, err } @@ -153,7 +155,8 @@ func (baid *Baidu) Search(query core.Query) ([]core.SearchResult, error) { // SearchImage executes a Baidu image search and returns normalized image // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (baid *Baidu) SearchImage(query core.Query) ([]core.SearchResult, error) { +func (baid *Baidu) SearchImage(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) baid.logger.Debug("Starting image search, query: %+v", query) searchResults := []core.SearchResult{} @@ -166,7 +169,7 @@ func (baid *Baidu) SearchImage(query core.Query) ([]core.SearchResult, error) { } // Get anti-crawler cookies first, then reload page - page, err := baid.Navigate(url) + page, err := baid.Navigate(ctx, url) if err != nil { return nil, err } diff --git a/baidu/search_integration_test.go b/baidu/search_integration_test.go index c609be5..dd3c8af 100644 --- a/baidu/search_integration_test.go +++ b/baidu/search_integration_test.go @@ -4,6 +4,7 @@ package baidu import ( + "context" "testing" "github.com/karust/openserp/core" @@ -18,7 +19,7 @@ func TestSearchBaidu(t *testing.T) { baid := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golang programming", Limit: 10} - results, err := baid.Search(query) + results, err := baid.Search(context.Background(), query) ithelper.HandleError(t, "baidu web search", err) if len(results) == 0 { @@ -39,7 +40,7 @@ func TestImageSearchBaidu(t *testing.T) { baid := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golden retriever puppy", Limit: 10} - results, err := baid.SearchImage(query) + results, err := baid.SearchImage(context.Background(), query) ithelper.HandleError(t, "baidu image search", err) if len(results) == 0 { diff --git a/baidu/search_raw.go b/baidu/search_raw.go index 7fd87d4..5d42b7e 100644 --- a/baidu/search_raw.go +++ b/baidu/search_raw.go @@ -1,6 +1,7 @@ package baidu import ( + "context" "net/http" "strings" @@ -10,13 +11,13 @@ import ( "github.com/sirupsen/logrus" ) -func baiduRequest(searchURL string, query core.Query) (*http.Response, error) { +func baiduRequest(ctx context.Context, searchURL string, query core.Query) (*http.Response, error) { baseClient, err := core.NewRawHTTPClient(query) if err != nil { return nil, err } - req, err := http.NewRequest("GET", searchURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { return nil, err } @@ -89,17 +90,20 @@ func baiduResultParser(response *http.Response) ([]core.SearchResult, error) { return results, err } -func Search(query core.Query) ([]core.SearchResult, error) { +func Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) + googleURL, err := BuildURL(query) if err != nil { return nil, err } logrus.Debugf("Baidu URL built: %s", googleURL) - res, err := baiduRequest(googleURL, query) + res, err := baiduRequest(ctx, googleURL, query) if err != nil { return nil, err } + defer res.Body.Close() logrus.Debugf("Baidu Raw response: code=%d", res.StatusCode) results, err := baiduResultParser(res) diff --git a/bing/search.go b/bing/search.go index b71ff41..d362907 100644 --- a/bing/search.go +++ b/bing/search.go @@ -1,6 +1,7 @@ package bing import ( + "context" "encoding/json" "errors" "fmt" @@ -81,15 +82,16 @@ func (bing *Bing) checkCaptcha(page *rod.Page) bool { return false } -func (bing *Bing) acceptCookies(page *rod.Page) { +func (bing *Bing) acceptCookies(ctx context.Context, page *rod.Page) error { consentBtn, err := page.Timeout(bing.Timeout / 10).Element("button#bnp_btn_accept") if err != nil { - return + return nil } if err := consentBtn.Click(proto.InputMouseButtonLeft, 1); err != nil { bing.logger.Debug("Cookie consent click failed: %v", err) } - time.Sleep(time.Millisecond * 500) + + return core.SleepContext(ctx, 500*time.Millisecond) } func (bing *Bing) close(page *rod.Page) { @@ -105,7 +107,8 @@ func (bing *Bing) close(page *rod.Page) { // Search executes a Bing web search and returns normalized search results. // It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (bing *Bing) Search(query core.Query) ([]core.SearchResult, error) { +func (bing *Bing) Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) bing.logger.Debug("Starting search, query: %+v", query) searchResults := []core.SearchResult{} @@ -115,7 +118,7 @@ func (bing *Bing) Search(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := bing.Navigate(url) + page, err := bing.Navigate(ctx, url) if err != nil { return nil, err } @@ -131,7 +134,9 @@ func (bing *Bing) Search(query core.Query) ([]core.SearchResult, error) { return nil, core.ErrCaptcha } - bing.acceptCookies(page) + if err := bing.acceptCookies(ctx, page); err != nil { + return nil, err + } if err := page.WaitLoad(); err != nil { bing.logger.Error("Post-consent page load wait failed: %s", err) return nil, core.ErrSearchTimeout @@ -260,7 +265,8 @@ type BingImageData struct { // SearchImage executes a Bing image search and returns normalized image // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (bing *Bing) SearchImage(query core.Query) ([]core.SearchResult, error) { +func (bing *Bing) SearchImage(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) bing.logger.Debug("Starting image search, query: %+v", query) searchResults := []core.SearchResult{} @@ -271,7 +277,7 @@ func (bing *Bing) SearchImage(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := bing.Navigate(url) + page, err := bing.Navigate(ctx, url) if err != nil { return nil, err } @@ -289,14 +295,18 @@ func (bing *Bing) SearchImage(query core.Query) ([]core.SearchResult, error) { } // Accept cookies if present - bing.acceptCookies(page) + if err := bing.acceptCookies(ctx, page); err != nil { + return nil, err + } // Wait for image results to load if err := page.WaitLoad(); err != nil { bing.logger.Error("Image results load wait failed: %s", err) return nil, core.ErrSearchTimeout } - time.Sleep(time.Second * 2) + if err := core.SleepContext(ctx, 2*time.Second); err != nil { + return nil, err + } // Find all image result containers using CSS selector imageContainers, err := page.Timeout(bing.Timeout).Elements("div.iuscp, div.isv") diff --git a/bing/search_integration_test.go b/bing/search_integration_test.go index 76530d9..a5bccea 100644 --- a/bing/search_integration_test.go +++ b/bing/search_integration_test.go @@ -4,6 +4,7 @@ package bing import ( + "context" "testing" "github.com/karust/openserp/core" @@ -18,7 +19,7 @@ func TestSearchBing(t *testing.T) { bing := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golang programming", Limit: 10} - results, err := bing.Search(query) + results, err := bing.Search(context.Background(), query) ithelper.HandleError(t, "bing web search", err) if len(results) == 0 { @@ -39,7 +40,7 @@ func TestImageSearchBing(t *testing.T) { bing := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golden retriever puppy", Limit: 10} - results, err := bing.SearchImage(query) + results, err := bing.SearchImage(context.Background(), query) ithelper.HandleError(t, "bing image search", err) if len(results) == 0 { diff --git a/cmd/root.go b/cmd/root.go index 42b6d38..834282e 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -14,7 +14,7 @@ import ( ) const ( - version = "0.6.3" + version = "0.6.4" defaultConfigFilename = "config" envPrefix = "OPENSERP" ) diff --git a/cmd/search.go b/cmd/search.go index 9082b5d..fc5bc30 100644 --- a/cmd/search.go +++ b/cmd/search.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "encoding/json" "fmt" "strings" @@ -129,19 +130,20 @@ func searchBrowser(engineType string, query core.Query, browserProxyURL string) return nil, fmt.Errorf("no %q search engine found", engineType) } - return engine.Search(query) + return engine.Search(context.Background(), query) } func searchRaw(engineType string, query core.Query) ([]core.SearchResult, error) { logrus.Warn("Browserless results are very inconsistent or may not even work!") + ctx := context.Background() switch strings.ToLower(engineType) { case "yandex": - return yandex.Search(query) + return yandex.Search(ctx, query) case "google": - return google.Search(query) + return google.Search(ctx, query) case "baidu": - return baidu.Search(query) + return baidu.Search(ctx, query) case "bing": logrus.Warn("Bing does not support raw HTTP requests mode. Please use browser mode instead.") return nil, fmt.Errorf("bing does not support raw requests mode") diff --git a/cmd/serve.go b/cmd/serve.go index 4ee3a2d..603eb5f 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "fmt" "strings" "sync" @@ -22,22 +23,22 @@ type rawEngine struct { name string } -func (r *rawEngine) Search(q core.Query) ([]core.SearchResult, error) { +func (r *rawEngine) Search(ctx context.Context, q core.Query) ([]core.SearchResult, error) { q.Insecure = config.Server.Insecure switch r.name { case "google": - return google.Search(q) + return google.Search(ctx, q) case "yandex": - return yandex.Search(q) + return yandex.Search(ctx, q) case "baidu": - return baidu.Search(q) + return baidu.Search(ctx, q) default: return nil, fmt.Errorf("unsupported engine: %s", r.name) } } -func (r *rawEngine) SearchImage(q core.Query) ([]core.SearchResult, error) { +func (r *rawEngine) SearchImage(_ context.Context, _ core.Query) ([]core.SearchResult, error) { return nil, fmt.Errorf("image search is not supported in raw mode for %s", r.name) } @@ -195,20 +196,20 @@ type pooledBrowserEngine struct { engines map[string]core.SearchEngine } -func (e *pooledBrowserEngine) Search(q core.Query) ([]core.SearchResult, error) { +func (e *pooledBrowserEngine) Search(ctx context.Context, q core.Query) ([]core.SearchResult, error) { engine, err := e.getOrCreate(q.ProxyURL) if err != nil { return nil, err } - return engine.Search(q) + return engine.Search(ctx, q) } -func (e *pooledBrowserEngine) SearchImage(q core.Query) ([]core.SearchResult, error) { +func (e *pooledBrowserEngine) SearchImage(ctx context.Context, q core.Query) ([]core.SearchResult, error) { engine, err := e.getOrCreate(q.ProxyURL) if err != nil { return nil, err } - return engine.SearchImage(q) + return engine.SearchImage(ctx, q) } func (e *pooledBrowserEngine) IsInitialized() bool { diff --git a/core/browser.go b/core/browser.go index 5845a36..10bef08 100644 --- a/core/browser.go +++ b/core/browser.go @@ -178,7 +178,12 @@ func (b *Browser) IsInitialized() bool { // Navigate connects to Chromium, creates a page, applies stealth/emulation and // proxy auth, then navigates to URL. It returns an initialized page ready for // selector queries, or an error when browser setup/navigation fails. -func (b *Browser) Navigate(URL string) (*rod.Page, error) { +func (b *Browser) Navigate(ctx context.Context, URL string) (*rod.Page, error) { + ctx = EnsureContext(ctx) + if err := ctx.Err(); err != nil { + return nil, err + } + logrus.Debug("Navigate to: ", URL) browser := rod.New().ControlURL(b.browserAddr).Timeout(b.Timeout) @@ -233,46 +238,48 @@ func (b *Browser) Navigate(URL string) (*rod.Page, error) { if err != nil { return nil, fmt.Errorf("create stealth page failed: %w", err) } - err = page.Emulate(devices.Device{ - AcceptLanguage: b.LanguageCode, - UserAgent: ua, - }) - if err != nil { - return nil, fmt.Errorf("emulate stealth page failed: %w", err) - } - } else { page, err = b.browser.Page(proto.TargetCreateTarget{URL: "about:blank"}) if err != nil { return nil, fmt.Errorf("create page failed: %w", err) } + } - err = page.Emulate(devices.Device{ - AcceptLanguage: b.LanguageCode, - UserAgent: ua, - }) - if err != nil { - return nil, fmt.Errorf("emulate page failed: %w", err) + // From here on, any error path must close the page to avoid leaking tabs + // when the caller context is canceled or navigation fails. + closeOnErr := func() { + if cerr := page.Close(); cerr != nil { + logrus.Debugf("Close page after navigate error failed: %v", cerr) } + } - err = proto.EmulationSetDeviceMetricsOverride{ + if err := page.Emulate(devices.Device{ + AcceptLanguage: b.LanguageCode, + UserAgent: ua, + }); err != nil { + closeOnErr() + return nil, fmt.Errorf("emulate page failed: %w", err) + } + + if !b.UseStealth { + if err := (proto.EmulationSetDeviceMetricsOverride{ Width: 1920, Height: 1080, DeviceScaleFactor: 1, Mobile: false, ScreenWidth: &[]int{1920}[0], ScreenHeight: &[]int{1080}[0], - }.Call(page) - if err != nil { + }).Call(page); err != nil { + closeOnErr() return nil, fmt.Errorf("set device metrics failed: %w", err) } } - //EnableCustomStealth(page) + page = page.Context(ctx) timedPage := page.Timeout(b.Timeout) - err = timedPage.Navigate(URL) - if err != nil { + if err := timedPage.Navigate(URL); err != nil { + closeOnErr() return nil, err } @@ -293,7 +300,10 @@ func (b *Browser) Navigate(URL string) (*rod.Page, error) { wait() } - time.Sleep(b.WaitLoadTime) + if err := SleepContext(ctx, b.WaitLoadTime); err != nil { + closeOnErr() + return nil, err + } return page, nil } diff --git a/core/browser_test.go b/core/browser_test.go index a5a9cc8..bd0b23f 100644 --- a/core/browser_test.go +++ b/core/browser_test.go @@ -4,6 +4,7 @@ package core import ( + "context" "encoding/json" "fmt" "os" @@ -34,7 +35,7 @@ func TestCreateBrowser(t *testing.T) { t.Fatalf("Error failed initializing browser: %s", err) } - page, err := browser.Navigate("about:blank") + page, err := browser.Navigate(context.Background(), "about:blank") if err != nil { t.Fatalf("navigate about:blank: %v", err) } @@ -133,7 +134,7 @@ func runSannysoftFingerprint(t *testing.T, useStealth bool) sannysoftRunSummary artifactPath := filepath.Join(botFingerprintArtifactDir, fmt.Sprintf("fingerprint_sannysoft_%s.png", label)) - page, err := browser.Navigate(sannysoftURL) + page, err := browser.Navigate(context.Background(), sannysoftURL) if err != nil { t.Fatalf("navigate to sannysoft (%s): %v", label, err) } diff --git a/core/context.go b/core/context.go new file mode 100644 index 0000000..62ce8bc --- /dev/null +++ b/core/context.go @@ -0,0 +1,38 @@ +package core + +import ( + "context" + "errors" + "time" +) + +// EnsureContext returns ctx when set; otherwise a non-nil placeholder context. +func EnsureContext(ctx context.Context) context.Context { + if ctx != nil { + return ctx + } + return context.TODO() +} + +// SleepContext blocks for d or until ctx is canceled. +func SleepContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + return nil + } + ctx = EnsureContext(ctx) + + timer := time.NewTimer(d) + defer timer.Stop() + + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// IsContextDone reports whether err is a cancellation/deadline error. +func IsContextDone(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) +} diff --git a/core/proxy_integration_test.go b/core/proxy_integration_test.go index bf21cfb..7812b70 100644 --- a/core/proxy_integration_test.go +++ b/core/proxy_integration_test.go @@ -78,7 +78,7 @@ func (e *proxyIntegrationEngine) GetRateLimiter() *rate.Limiter { return e.limiter } -func (e *proxyIntegrationEngine) Search(q Query) ([]SearchResult, error) { +func (e *proxyIntegrationEngine) Search(_ context.Context, q Query) ([]SearchResult, error) { e.proxies = append(e.proxies, q.ProxyURL) body, err := fetchViaRawProxy(q.ProxyURL, q.Insecure, e.targetURL) @@ -94,7 +94,8 @@ func (e *proxyIntegrationEngine) Search(q Query) ([]SearchResult, error) { }}, nil } -func (e *proxyIntegrationEngine) SearchImage(q Query) ([]SearchResult, error) { +func (e *proxyIntegrationEngine) SearchImage(_ context.Context, q Query) ([]SearchResult, error) { + _ = q return nil, ErrSearchTimeout } diff --git a/core/resilient.go b/core/resilient.go index 8fee752..0ac3757 100644 --- a/core/resilient.go +++ b/core/resilient.go @@ -74,8 +74,8 @@ func NewResilientSearcher(engines []SearchEngine, cfg ResilientConfig) *Resilien } // SearchPrimary keeps dedicated endpoints engine-pure (no fallback). -func (rs *ResilientSearcher) SearchPrimary(primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { - results, proxyMeta, err := rs.searchWithProtection(primaryEngine, q, false) +func (rs *ResilientSearcher) SearchPrimary(ctx context.Context, primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { + results, proxyMeta, err := rs.searchWithProtection(ctx, primaryEngine, q, false) if err != nil { return nil, primaryEngine.Name(), proxyMeta, err } @@ -83,27 +83,32 @@ func (rs *ResilientSearcher) SearchPrimary(primaryEngine SearchEngine, q Query) } // SearchWithFallback retries primary and then tries other initialized engines. -func (rs *ResilientSearcher) SearchWithFallback(primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { - return rs.searchWithFallback(primaryEngine, q, false) +func (rs *ResilientSearcher) SearchWithFallback(ctx context.Context, primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { + return rs.searchWithFallback(ctx, primaryEngine, q, false) } -func (rs *ResilientSearcher) SearchImagePrimary(primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { - results, proxyMeta, err := rs.searchWithProtection(primaryEngine, q, true) +func (rs *ResilientSearcher) SearchImagePrimary(ctx context.Context, primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { + results, proxyMeta, err := rs.searchWithProtection(ctx, primaryEngine, q, true) if err != nil { return nil, primaryEngine.Name(), proxyMeta, err } return results, primaryEngine.Name(), proxyMeta, nil } -func (rs *ResilientSearcher) SearchImageWithFallback(primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { - return rs.searchWithFallback(primaryEngine, q, true) +func (rs *ResilientSearcher) SearchImageWithFallback(ctx context.Context, primaryEngine SearchEngine, q Query) ([]SearchResult, string, ProxyExecutionMeta, error) { + return rs.searchWithFallback(ctx, primaryEngine, q, true) } -func (rs *ResilientSearcher) searchWithFallback(primaryEngine SearchEngine, q Query, isImage bool) ([]SearchResult, string, ProxyExecutionMeta, error) { - results, proxyMeta, err := rs.searchWithProtection(primaryEngine, q, isImage) +func (rs *ResilientSearcher) searchWithFallback(ctx context.Context, primaryEngine SearchEngine, q Query, isImage bool) ([]SearchResult, string, ProxyExecutionMeta, error) { + ctx = EnsureContext(ctx) + + results, proxyMeta, err := rs.searchWithProtection(ctx, primaryEngine, q, isImage) if err == nil { return results, primaryEngine.Name(), proxyMeta, nil } + if ctx.Err() != nil { + return nil, primaryEngine.Name(), proxyMeta, ctx.Err() + } if errors.Is(err, ErrProxyUnavailable) { logrus.Warnf("[Resilient] Primary engine %s proxy policy failed closed: %s", primaryEngine.Name(), err) return nil, primaryEngine.Name(), proxyMeta, err @@ -118,11 +123,14 @@ func (rs *ResilientSearcher) searchWithFallback(primaryEngine SearchEngine, q Qu logrus.Warnf("[Resilient] Primary engine %s %s: %s. Trying fallback engines...", primaryEngine.Name(), action, err) for _, fallbackEngine := range rs.engines { + if ctx.Err() != nil { + return nil, primaryEngine.Name(), proxyMeta, ctx.Err() + } if fallbackEngine.Name() == primaryEngine.Name() || !fallbackEngine.IsInitialized() { continue } - results, fallbackMeta, fallbackErr := rs.searchWithProtection(fallbackEngine, q, isImage) + results, fallbackMeta, fallbackErr := rs.searchWithProtection(ctx, fallbackEngine, q, isImage) if fallbackErr == nil { logrus.Infof("[Resilient] "+successMessage, fallbackEngine.Name(), len(results)) return results, fallbackEngine.Name(), fallbackMeta, nil @@ -133,7 +141,12 @@ func (rs *ResilientSearcher) searchWithFallback(primaryEngine SearchEngine, q Qu return nil, primaryEngine.Name(), proxyMeta, ErrAllEnginesFailed } -func (rs *ResilientSearcher) searchWithProtection(engine SearchEngine, q Query, isImage bool) ([]SearchResult, ProxyExecutionMeta, error) { +func (rs *ResilientSearcher) searchWithProtection(ctx context.Context, engine SearchEngine, q Query, isImage bool) ([]SearchResult, ProxyExecutionMeta, error) { + ctx = EnsureContext(ctx) + + if ctx.Err() != nil { + return nil, ProxyExecutionMeta{}, ctx.Err() + } cb := rs.cbManager.Get(engine.Name()) if !cb.AllowRequest() { return nil, ProxyExecutionMeta{}, ErrCircuitOpen @@ -142,10 +155,10 @@ func (rs *ResilientSearcher) searchWithProtection(engine SearchEngine, q Query, policy := rs.effectivePolicyForQuery(engine.Name(), q) attemptMeta := rs.baseProxyMeta(policy) - result := RetryableSearch(rs.retryCfg, engine.Name(), func() ([]SearchResult, error) { + result := RetryableSearch(ctx, rs.retryCfg, engine.Name(), func(callCtx context.Context) ([]SearchResult, error) { limiter := engine.GetRateLimiter() if limiter != nil { - if err := limiter.Wait(context.Background()); err != nil { + if err := limiter.Wait(callCtx); err != nil { return nil, err } } @@ -174,9 +187,9 @@ func (rs *ResilientSearcher) searchWithProtection(engine SearchEngine, q Query, err error ) if isImage { - results, err = engine.SearchImage(attemptQuery) + results, err = engine.SearchImage(callCtx, attemptQuery) } else { - results, err = engine.Search(attemptQuery) + results, err = engine.Search(callCtx, attemptQuery) } if reportToRegistry { @@ -198,12 +211,17 @@ func (rs *ResilientSearcher) searchWithProtection(engine SearchEngine, q Query, } // SearchAllParallel applies retry/circuit protections per engine for mega search. -func (rs *ResilientSearcher) SearchAllParallel(q Query, engines []SearchEngine) []MegaSearchResult { +func (rs *ResilientSearcher) SearchAllParallel(ctx context.Context, q Query, engines []SearchEngine) []MegaSearchResult { + ctx = EnsureContext(ctx) + var wg sync.WaitGroup var mu sync.Mutex var allResults []MegaSearchResult for _, engine := range engines { + if ctx.Err() != nil { + break + } if !engine.IsInitialized() { continue } @@ -216,7 +234,7 @@ func (rs *ResilientSearcher) SearchAllParallel(q Query, engines []SearchEngine) go func(eng SearchEngine) { defer wg.Done() - results, _, err := rs.searchWithProtection(eng, q, false) + results, _, err := rs.searchWithProtection(ctx, eng, q, false) if err != nil { return } @@ -236,12 +254,17 @@ func (rs *ResilientSearcher) SearchAllParallel(q Query, engines []SearchEngine) return allResults } -func (rs *ResilientSearcher) SearchAllImageParallel(q Query, engines []SearchEngine) []MegaSearchResult { +func (rs *ResilientSearcher) SearchAllImageParallel(ctx context.Context, q Query, engines []SearchEngine) []MegaSearchResult { + ctx = EnsureContext(ctx) + var wg sync.WaitGroup var mu sync.Mutex var allResults []MegaSearchResult for _, engine := range engines { + if ctx.Err() != nil { + break + } if !engine.IsInitialized() { continue } @@ -254,7 +277,7 @@ func (rs *ResilientSearcher) SearchAllImageParallel(q Query, engines []SearchEng go func(eng SearchEngine) { defer wg.Done() - results, _, err := rs.searchWithProtection(eng, q, true) + results, _, err := rs.searchWithProtection(ctx, eng, q, true) if err != nil { return } diff --git a/core/resilient_context_test.go b/core/resilient_context_test.go new file mode 100644 index 0000000..0a2e057 --- /dev/null +++ b/core/resilient_context_test.go @@ -0,0 +1,75 @@ +package core + +import ( + "context" + "errors" + "testing" + "time" + + "golang.org/x/time/rate" +) + +type blockingContextEngine struct { + started chan struct{} +} + +func (e *blockingContextEngine) Name() string { + return "blocking" +} + +func (e *blockingContextEngine) IsInitialized() bool { + return true +} + +func (e *blockingContextEngine) GetRateLimiter() *rate.Limiter { + return nil +} + +func (e *blockingContextEngine) Search(ctx context.Context, q Query) ([]SearchResult, error) { + _ = q + close(e.started) + <-ctx.Done() + return nil, ctx.Err() +} + +func (e *blockingContextEngine) SearchImage(ctx context.Context, q Query) ([]SearchResult, error) { + _ = q + <-ctx.Done() + return nil, ctx.Err() +} + +func TestResilientSearchPrimary_CancelledContextStopsWithin100ms(t *testing.T) { + engine := &blockingContextEngine{started: make(chan struct{})} + cfg := DefaultResilientConfig() + cfg.Retry.MaxRetries = 2 + rs := NewResilientSearcher([]SearchEngine{engine}, cfg) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + + go func() { + _, _, _, err := rs.SearchPrimary(ctx, engine, Query{Text: "cancel-me"}) + done <- err + }() + + select { + case <-engine.started: + case <-time.After(2 * time.Second): + t.Fatal("search did not start") + } + + start := time.Now() + cancel() + + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if elapsed := time.Since(start); elapsed > 100*time.Millisecond { + t.Fatalf("expected cancellation within 100ms, got %s", elapsed) + } + case <-time.After(2 * time.Second): + t.Fatal("search did not stop after context cancellation") + } +} diff --git a/core/retry.go b/core/retry.go index bbf48d9..f1e1fce 100644 --- a/core/retry.go +++ b/core/retry.go @@ -1,9 +1,11 @@ package core import ( + "context" "errors" "fmt" "math" + "math/rand" "time" "github.com/sirupsen/logrus" @@ -35,20 +37,35 @@ type RetryResult struct { // RetryableSearch executes searchFn with exponential backoff retries. // CAPTCHA and proxy-unavailable errors are not retried. -func RetryableSearch(cfg RetryConfig, engineName string, searchFn func() ([]SearchResult, error)) RetryResult { +func RetryableSearch(ctx context.Context, cfg RetryConfig, engineName string, searchFn func(context.Context) ([]SearchResult, error)) RetryResult { + ctx = EnsureContext(ctx) if cfg.BackoffFactor <= 0 { cfg.BackoffFactor = 2.0 } var lastErr error for attempt := 0; attempt <= cfg.MaxRetries; attempt++ { + if err := ctx.Err(); err != nil { + return RetryResult{ + Err: err, + Attempts: attempt, + Engine: engineName, + } + } + if attempt > 0 { backoff := calculateBackoff(cfg, attempt) logrus.Warnf("[%s] Retry attempt %d/%d after %s", engineName, attempt, cfg.MaxRetries, backoff) - time.Sleep(backoff) + if err := SleepContext(ctx, backoff); err != nil { + return RetryResult{ + Err: err, + Attempts: attempt, + Engine: engineName, + } + } } - results, err := searchFn() + results, err := searchFn(ctx) if err == nil { if attempt > 0 { logrus.Infof("[%s] Succeeded on retry attempt %d", engineName, attempt) @@ -77,6 +94,14 @@ func RetryableSearch(cfg RetryConfig, engineName string, searchFn func() ([]Sear Engine: engineName, } } + if IsContextDone(err) { + logrus.Warnf("[%s] Context canceled/deadline exceeded, skipping retries", engineName) + return RetryResult{ + Err: err, + Attempts: attempt + 1, + Engine: engineName, + } + } logrus.Warnf("[%s] Attempt %d failed: %s", engineName, attempt+1, err) } @@ -96,5 +121,9 @@ func calculateBackoff(cfg RetryConfig, attempt int) time.Duration { if backoff < 0 { backoff = 0 } + backoff = backoff * (0.5 + rand.Float64()) + if backoff > float64(cfg.MaxBackoff) { + backoff = float64(cfg.MaxBackoff) + } return time.Duration(backoff) } diff --git a/core/retry_test.go b/core/retry_test.go index ca3faf6..7aded6e 100644 --- a/core/retry_test.go +++ b/core/retry_test.go @@ -1,6 +1,7 @@ package core import ( + "context" "errors" "testing" "time" @@ -10,8 +11,11 @@ func TestRetryableSearch_SuccessOnFirstAttempt(t *testing.T) { cfg := RetryConfig{MaxRetries: 3, InitialBackoff: 10 * time.Millisecond, MaxBackoff: 100 * time.Millisecond, BackoffFactor: 2.0} calls := 0 - result := RetryableSearch(cfg, "test", func() ([]SearchResult, error) { + result := RetryableSearch(context.Background(), cfg, "test", func(ctx context.Context) ([]SearchResult, error) { calls++ + if ctx == nil { + t.Fatal("expected non-nil context") + } return []SearchResult{{Title: "result1"}}, nil }) @@ -30,7 +34,7 @@ func TestRetryableSearch_AllAttemptsFail(t *testing.T) { cfg := RetryConfig{MaxRetries: 2, InitialBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond, BackoffFactor: 2.0} calls := 0 - result := RetryableSearch(cfg, "test", func() ([]SearchResult, error) { + result := RetryableSearch(context.Background(), cfg, "test", func(context.Context) ([]SearchResult, error) { calls++ return nil, errors.New("persistent failure") }) @@ -50,7 +54,7 @@ func TestRetryableSearch_CaptchaNotRetried(t *testing.T) { cfg := RetryConfig{MaxRetries: 3, InitialBackoff: 10 * time.Millisecond, MaxBackoff: 100 * time.Millisecond, BackoffFactor: 2.0} calls := 0 - result := RetryableSearch(cfg, "test", func() ([]SearchResult, error) { + result := RetryableSearch(context.Background(), cfg, "test", func(context.Context) ([]SearchResult, error) { calls++ return nil, ErrCaptcha }) @@ -63,24 +67,54 @@ func TestRetryableSearch_CaptchaNotRetried(t *testing.T) { } } +func TestRetryableSearch_ContextCanceledStopsBackoffImmediately(t *testing.T) { + cfg := RetryConfig{ + MaxRetries: 3, + InitialBackoff: time.Second, + MaxBackoff: time.Second, + BackoffFactor: 2.0, + } + ctx, cancel := context.WithCancel(context.Background()) + calls := 0 + + start := time.Now() + result := RetryableSearch(ctx, cfg, "test", func(context.Context) ([]SearchResult, error) { + calls++ + cancel() + return nil, errors.New("trigger retry") + }) + elapsed := time.Since(start) + + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("expected context.Canceled, got: %v", result.Err) + } + if calls != 1 { + t.Fatalf("expected a single attempt before cancel, got %d", calls) + } + if elapsed > 100*time.Millisecond { + t.Fatalf("expected fast cancel, got %s", elapsed) + } +} + func TestCalculateBackoff(t *testing.T) { cfg := RetryConfig{InitialBackoff: 1 * time.Second, MaxBackoff: 10 * time.Second, BackoffFactor: 2.0} tests := []struct { - attempt int - expected time.Duration + attempt int + min time.Duration + max time.Duration }{ - {1, 1 * time.Second}, - {2, 2 * time.Second}, - {3, 4 * time.Second}, - {4, 8 * time.Second}, - {5, 10 * time.Second}, + {1, 500 * time.Millisecond, 1500 * time.Millisecond}, + {2, 1 * time.Second, 3 * time.Second}, + {3, 2 * time.Second, 6 * time.Second}, + {4, 4 * time.Second, 10 * time.Second}, + {5, 5 * time.Second, 10 * time.Second}, } for _, tt := range tests { got := calculateBackoff(cfg, tt.attempt) - if got != tt.expected { - t.Errorf("attempt %d: expected %s, got %s", tt.attempt, tt.expected, got) + if got < tt.min || got > tt.max { + t.Errorf("attempt %d: expected range [%s,%s], got %s", tt.attempt, tt.min, tt.max, got) } } } diff --git a/core/server.go b/core/server.go index 2eb9c8d..d4373e5 100644 --- a/core/server.go +++ b/core/server.go @@ -1,6 +1,7 @@ package core import ( + "context" "encoding/json" "errors" "fmt" @@ -22,9 +23,9 @@ type SearchEngine interface { // Search runs a web search request and returns normalized results. // Implementations should return sentinel errors such as ErrCaptcha and // ErrSearchTimeout for policy-aware handling. - Search(Query) ([]SearchResult, error) + Search(context.Context, Query) ([]SearchResult, error) // SearchImage runs an image search request and returns normalized results. - SearchImage(Query) ([]SearchResult, error) + SearchImage(context.Context, Query) ([]SearchResult, error) // IsInitialized reports whether the engine is ready to serve requests. IsInitialized() bool // Name returns a stable engine identifier used in routes and telemetry. @@ -176,15 +177,15 @@ func (s *Server) handleDedicatedEndpoint(c *fiber.Ctx, engine SearchEngine, isIm if isImage { if s.opts.AllowEndpointFallback { - res, usedEngine, proxyMeta, searchErr = s.resilient.SearchImageWithFallback(engine, q) + res, usedEngine, proxyMeta, searchErr = s.resilient.SearchImageWithFallback(c.UserContext(), engine, q) } else { - res, usedEngine, proxyMeta, searchErr = s.resilient.SearchImagePrimary(engine, q) + res, usedEngine, proxyMeta, searchErr = s.resilient.SearchImagePrimary(c.UserContext(), engine, q) } } else { if s.opts.AllowEndpointFallback { - res, usedEngine, proxyMeta, searchErr = s.resilient.SearchWithFallback(engine, q) + res, usedEngine, proxyMeta, searchErr = s.resilient.SearchWithFallback(c.UserContext(), engine, q) } else { - res, usedEngine, proxyMeta, searchErr = s.resilient.SearchPrimary(engine, q) + res, usedEngine, proxyMeta, searchErr = s.resilient.SearchPrimary(c.UserContext(), engine, q) } } s.applyProxyHeaders(c, proxyMeta) @@ -350,7 +351,7 @@ func (s *Server) handleMegaImage(c *fiber.Ctx) error { return s.handleMegaEndpoint(c, "image", s.resilient.SearchAllImageParallel) } -func (s *Server) handleMegaEndpoint(c *fiber.Ctx, action string, run func(Query, []SearchEngine) []MegaSearchResult) error { +func (s *Server) handleMegaEndpoint(c *fiber.Ctx, action string, run func(context.Context, Query, []SearchEngine) []MegaSearchResult) error { q := Query{} if err := q.InitFromContext(c); err != nil { logrus.Errorf("Error while setting mega %s query: %s", action, err) @@ -387,7 +388,7 @@ func (s *Server) handleMegaEndpoint(c *fiber.Ctx, action string, run func(Query, return err } - results := run(q, enginesToUse) + results := run(c.UserContext(), q, enginesToUse) dedupedResults := s.deduplicateMegaResults(results) if s.cache != nil { diff --git a/core/server_test.go b/core/server_test.go index 6aa51b2..65e4418 100644 --- a/core/server_test.go +++ b/core/server_test.go @@ -1,6 +1,7 @@ package core import ( + "context" "encoding/json" "errors" "io" @@ -18,8 +19,8 @@ type engineMock struct { name string initialized bool limiter *rate.Limiter - searchFn func(Query) ([]SearchResult, error) - imageFn func(Query) ([]SearchResult, error) + searchFn func(context.Context, Query) ([]SearchResult, error) + imageFn func(context.Context, Query) ([]SearchResult, error) mu sync.Mutex searchCalls int @@ -32,22 +33,22 @@ func (e *engineMock) IsInitialized() bool { } func (e *engineMock) GetRateLimiter() *rate.Limiter { return e.limiter } -func (e *engineMock) Search(q Query) ([]SearchResult, error) { +func (e *engineMock) Search(ctx context.Context, q Query) ([]SearchResult, error) { e.mu.Lock() e.searchCalls++ e.mu.Unlock() if e.searchFn != nil { - return e.searchFn(q) + return e.searchFn(ctx, q) } return []SearchResult{{Rank: 1, URL: "https://example.com/" + e.name, Title: e.name}}, nil } -func (e *engineMock) SearchImage(q Query) ([]SearchResult, error) { +func (e *engineMock) SearchImage(ctx context.Context, q Query) ([]SearchResult, error) { e.mu.Lock() e.imageCalls++ e.mu.Unlock() if e.imageFn != nil { - return e.imageFn(q) + return e.imageFn(ctx, q) } return []SearchResult{{Rank: 1, URL: "https://img.example.com/" + e.name, Title: e.name}}, nil } @@ -366,7 +367,7 @@ func TestDedicatedEndpointNoFallbackByDefault(t *testing.T) { primary := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return nil, errors.New("primary failed") }, } @@ -453,7 +454,7 @@ func TestDedicatedEndpointFallbackBypassesCache(t *testing.T) { primary := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return nil, errors.New("primary failed") }, } @@ -524,14 +525,14 @@ func TestMegaSearchCachesWholeQueryWithEngineOrderNormalization(t *testing.T) { google := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return []SearchResult{{Rank: 1, URL: "https://example.com/shared", Title: "shared"}}, nil }, } yandex := &engineMock{ name: "yandex", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return []SearchResult{{Rank: 1, URL: "https://example.com/shared", Title: "shared"}}, nil }, } @@ -567,14 +568,14 @@ func TestMegaImageCachesWholeQueryWithEngineOrderNormalization(t *testing.T) { google := &engineMock{ name: "google", initialized: true, - imageFn: func(q Query) ([]SearchResult, error) { + imageFn: func(_ context.Context, q Query) ([]SearchResult, error) { return []SearchResult{{Rank: 1, URL: "https://img.example.com/shared", Title: "shared"}}, nil }, } yandex := &engineMock{ name: "yandex", initialized: true, - imageFn: func(q Query) ([]SearchResult, error) { + imageFn: func(_ context.Context, q Query) ([]SearchResult, error) { return []SearchResult{{Rank: 1, URL: "https://img.example.com/shared", Title: "shared"}}, nil }, } @@ -628,7 +629,7 @@ func TestMegaSearchCachesForHealthySubsetWhenOneCircuitIsOpen(t *testing.T) { bing := &engineMock{ name: "bing", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return nil, errors.New("bing failed") }, } @@ -666,7 +667,7 @@ func TestResilienceStatsContainsRetryInWhenCircuitOpen(t *testing.T) { primary := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return nil, errors.New("forced failure") }, } @@ -795,7 +796,7 @@ func TestResilientRawProxyPoolRotatesOnRetry(t *testing.T) { engine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { attemptedProxies = append(attemptedProxies, q.ProxyURL) if q.ProxyURL == "http://bad-proxy:8080" { return nil, errors.New("proxy failed") @@ -898,7 +899,7 @@ func TestGlobalProxyForcesAllEnginesRaw(t *testing.T) { googleEngine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { googleProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -906,7 +907,7 @@ func TestGlobalProxyForcesAllEnginesRaw(t *testing.T) { yandexEngine := &engineMock{ name: "yandex", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { yandexProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/yandex", Title: "yandex"}}, nil }, @@ -944,7 +945,7 @@ func TestRequestProxyOverrideDirectBeatsGlobal(t *testing.T) { engine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { googleProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -979,7 +980,7 @@ func TestRequestProxyOverrideTagBeatsGlobal(t *testing.T) { engine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { googleProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -1017,7 +1018,7 @@ func TestBrowserProxyPoolRotatesPerRequest(t *testing.T) { engine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { attemptedProxies = append(attemptedProxies, q.ProxyURL) return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -1079,7 +1080,7 @@ func TestMegaProxyOverrideHeaderBeatsGlobal(t *testing.T) { engine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { googleProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -1137,7 +1138,7 @@ func TestEngineOverrideProxyBehaviorRaw(t *testing.T) { googleEngine := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { googleProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/google", Title: "google"}}, nil }, @@ -1145,7 +1146,7 @@ func TestEngineOverrideProxyBehaviorRaw(t *testing.T) { yandexEngine := &engineMock{ name: "yandex", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { yandexProxy = q.ProxyURL return []SearchResult{{Rank: 1, URL: "https://example.com/yandex", Title: "yandex"}}, nil }, @@ -1183,7 +1184,7 @@ func TestRetryAppliesRateLimiterOnEachAttempt(t *testing.T) { name: "google", initialized: true, limiter: rate.NewLimiter(rate.Every(120*time.Millisecond), 1), - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { return nil, errors.New("always fail") }, } @@ -1236,7 +1237,7 @@ func TestCacheStatsReflectActivity(t *testing.T) { primary := &engineMock{ name: "google", initialized: true, - searchFn: func(q Query) ([]SearchResult, error) { + searchFn: func(_ context.Context, q Query) ([]SearchResult, error) { if q.Text == "fallback" { return nil, errors.New("fallback path") } diff --git a/duckduckgo/search.go b/duckduckgo/search.go index eb66b78..d868ffd 100644 --- a/duckduckgo/search.go +++ b/duckduckgo/search.go @@ -1,6 +1,7 @@ package duckduckgo import ( + "context" "fmt" "strings" "time" @@ -169,7 +170,8 @@ func (ddg *DuckDuckGo) parseResults(results rod.Elements, pageNum int) []core.Se // Search executes a DuckDuckGo web search and returns normalized search // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (ddg *DuckDuckGo) Search(query core.Query) ([]core.SearchResult, error) { +func (ddg *DuckDuckGo) Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) ddg.logger.Debug("Starting search, query: %+v", query) allResults := []core.SearchResult{} @@ -181,7 +183,7 @@ func (ddg *DuckDuckGo) Search(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := ddg.Navigate(url) + page, err := ddg.Navigate(ctx, url) if err != nil { return nil, err } @@ -256,7 +258,9 @@ func (ddg *DuckDuckGo) Search(query core.Query) ([]core.SearchResult, error) { break } - time.Sleep(ddg.pageSleep) + if err := core.SleepContext(ctx, ddg.pageSleep); err != nil { + return nil, err + } } // Deduplicate results @@ -273,7 +277,8 @@ func (ddg *DuckDuckGo) Search(query core.Query) ([]core.SearchResult, error) { // SearchImage executes a DuckDuckGo image search and returns normalized image // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (ddg *DuckDuckGo) SearchImage(query core.Query) ([]core.SearchResult, error) { +func (ddg *DuckDuckGo) SearchImage(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) ddg.logger.Debug("Starting image search, query: %+v", query) searchResults := []core.SearchResult{} @@ -283,7 +288,7 @@ func (ddg *DuckDuckGo) SearchImage(query core.Query) ([]core.SearchResult, error return nil, err } - page, err := ddg.Navigate(url) + page, err := ddg.Navigate(ctx, url) if err != nil { return nil, err } @@ -297,7 +302,9 @@ func (ddg *DuckDuckGo) SearchImage(query core.Query) ([]core.SearchResult, error ddg.logger.Error("Wait load failed: %s", err) return searchResults, core.ErrSearchTimeout } - time.Sleep(time.Second * 2) // Give time for images to load + if err := core.SleepContext(ctx, 2*time.Second); err != nil { + return searchResults, err + } // Try multiple selectors for DuckDuckGo image results var searchRes *rod.SearchResult diff --git a/duckduckgo/search_integration_test.go b/duckduckgo/search_integration_test.go index 4c4b0d3..2261cbd 100644 --- a/duckduckgo/search_integration_test.go +++ b/duckduckgo/search_integration_test.go @@ -4,6 +4,7 @@ package duckduckgo import ( + "context" "testing" "github.com/karust/openserp/core" @@ -18,7 +19,7 @@ func TestSearchDuckDuckGo(t *testing.T) { ddg := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "wikipedia", Limit: 10} - results, err := ddg.Search(query) + results, err := ddg.Search(context.Background(), query) ithelper.HandleError(t, "duckduckgo web search", err) if len(results) == 0 { @@ -39,7 +40,7 @@ func TestImageSearchDuckDuckGo(t *testing.T) { ddg := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golden retriever puppy", Limit: 10} - results, err := ddg.SearchImage(query) + results, err := ddg.SearchImage(context.Background(), query) ithelper.HandleError(t, "duckduckgo image search", err) if len(results) == 0 { diff --git a/google/search.go b/google/search.go index 97861b9..952ed64 100644 --- a/google/search.go +++ b/google/search.go @@ -1,6 +1,7 @@ package google import ( + "context" "errors" "fmt" "regexp" @@ -163,7 +164,8 @@ func (gogl *Google) acceptCookies(page *rod.Page) { // Search executes a Google web search and returns normalized search results. // It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (gogl *Google) Search(query core.Query) ([]core.SearchResult, error) { +func (gogl *Google) Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) gogl.logger.Debug("Starting search, query: %+v", query) searchResults := []core.SearchResult{} @@ -173,7 +175,7 @@ func (gogl *Google) Search(query core.Query) ([]core.SearchResult, error) { if err != nil { return nil, err } - page, err := gogl.Navigate(url) + page, err := gogl.Navigate(ctx, url) if err != nil { return nil, err } @@ -278,7 +280,9 @@ func (gogl *Google) Search(query core.Query) ([]core.SearchResult, error) { } //answ.Page().WaitRepaint() } - time.Sleep(time.Millisecond * 2000) + if err := core.SleepContext(ctx, 2*time.Second); err != nil { + return nil, err + } for i, answ := range answers { answerText := strings.Split(answ.MustText(), "\n") @@ -374,7 +378,8 @@ func (gogl *Google) Search(query core.Query) ([]core.SearchResult, error) { // SearchImage executes a Google image search and returns normalized image // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (gogl *Google) SearchImage(query core.Query) ([]core.SearchResult, error) { +func (gogl *Google) SearchImage(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) gogl.logger.Debug("Starting image search, query: %+v", query) searchResultsMap := map[string]core.SearchResult{} @@ -383,7 +388,7 @@ func (gogl *Google) SearchImage(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := gogl.Navigate(url) + page, err := gogl.Navigate(ctx, url) if err != nil { return nil, err } diff --git a/google/search_integration_test.go b/google/search_integration_test.go index 912f3ad..f3e0613 100644 --- a/google/search_integration_test.go +++ b/google/search_integration_test.go @@ -4,6 +4,7 @@ package google import ( + "context" "testing" "github.com/karust/openserp/core" @@ -18,7 +19,7 @@ func TestSearchGoogle(t *testing.T) { gogl := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golang programming", Limit: 10} - results, err := gogl.Search(query) + results, err := gogl.Search(context.Background(), query) ithelper.HandleError(t, "google web search", err) if len(results) == 0 { @@ -39,7 +40,7 @@ func TestImageSearchGoogle(t *testing.T) { gogl := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "golden retriever puppy", Limit: 10} - results, err := gogl.SearchImage(query) + results, err := gogl.SearchImage(context.Background(), query) ithelper.HandleError(t, "google image search", err) if len(results) == 0 { diff --git a/google/search_raw.go b/google/search_raw.go index ed65730..dcf8c30 100644 --- a/google/search_raw.go +++ b/google/search_raw.go @@ -1,6 +1,7 @@ package google import ( + "context" "net/http" "strings" @@ -10,13 +11,13 @@ import ( "github.com/sirupsen/logrus" ) -func googleRequest(searchURL string, query core.Query) (*http.Response, error) { +func googleRequest(ctx context.Context, searchURL string, query core.Query) (*http.Response, error) { baseClient, err := core.NewRawHTTPClient(query) if err != nil { return nil, err } - req, err := http.NewRequest("GET", searchURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { return nil, err } @@ -100,17 +101,20 @@ func googleResultParser(response *http.Response) ([]core.SearchResult, error) { return core.DeduplicateResults(results), err } -func Search(query core.Query) ([]core.SearchResult, error) { +func Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) + googleURL, err := BuildURL(query) if err != nil { return nil, err } logrus.Debugf("Google URL built: %s", googleURL) - res, err := googleRequest(googleURL, query) + res, err := googleRequest(ctx, googleURL, query) if err != nil { return nil, err } + defer res.Body.Close() logrus.Debugf("Google Raw response: code=%d", res.StatusCode) results, err := googleResultParser(res) diff --git a/yandex/search.go b/yandex/search.go index ac2eac6..76a389b 100644 --- a/yandex/search.go +++ b/yandex/search.go @@ -1,6 +1,7 @@ package yandex import ( + "context" "encoding/json" "fmt" "sort" @@ -121,7 +122,8 @@ func (yand *Yandex) parseResults(results rod.Elements, pageNum int) []core.Searc // Search executes a Yandex web search and returns normalized search results. // It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (yand *Yandex) Search(query core.Query) ([]core.SearchResult, error) { +func (yand *Yandex) Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) yand.logger.Debug("Starting search, query: %+v", query) if query.Start < 0 { return nil, fmt.Errorf("incorrect start provided") @@ -141,7 +143,7 @@ func (yand *Yandex) Search(query core.Query) ([]core.SearchResult, error) { return nil, err } - page, err := yand.Navigate(url) + page, err := yand.Navigate(ctx, url) if err != nil { return nil, err } @@ -193,7 +195,9 @@ func (yand *Yandex) Search(query core.Query) ([]core.SearchResult, error) { } } - time.Sleep(yand.pageSleep) + if err := core.SleepContext(ctx, yand.pageSleep); err != nil { + return nil, err + } } yand.logger.Info("Search completed: %d results", len(allResults)) @@ -202,7 +206,8 @@ func (yand *Yandex) Search(query core.Query) ([]core.SearchResult, error) { // SearchImage executes a Yandex image search and returns normalized image // results. It may return core.ErrCaptcha or core.ErrSearchTimeout. -func (yand *Yandex) SearchImage(query core.Query) ([]core.SearchResult, error) { +func (yand *Yandex) SearchImage(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) yand.logger.Debug("Starting image search, query: %+v", query) searchResults := []core.SearchResult{} @@ -215,7 +220,7 @@ func (yand *Yandex) SearchImage(query core.Query) ([]core.SearchResult, error) { } searchPage += 1 - page, err := yand.Navigate(url) + page, err := yand.Navigate(ctx, url) if err != nil { return nil, err } diff --git a/yandex/search_raw.go b/yandex/search_raw.go index 38d9a31..93fdd3a 100644 --- a/yandex/search_raw.go +++ b/yandex/search_raw.go @@ -1,6 +1,7 @@ package yandex import ( + "context" "net/http" "strings" @@ -10,13 +11,13 @@ import ( "github.com/sirupsen/logrus" ) -func yandexRequest(searchURL string, query core.Query) (*http.Response, error) { +func yandexRequest(ctx context.Context, searchURL string, query core.Query) (*http.Response, error) { baseClient, err := core.NewRawHTTPClient(query) if err != nil { return nil, err } - req, err := http.NewRequest("GET", searchURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { return nil, err } @@ -88,7 +89,9 @@ func yandexResultParser(response *http.Response) ([]core.SearchResult, error) { return core.DeduplicateResults(results), err } -func Search(query core.Query) ([]core.SearchResult, error) { +func Search(ctx context.Context, query core.Query) ([]core.SearchResult, error) { + ctx = core.EnsureContext(ctx) + startPage, skipOnFirstPage, err := core.ComputePagination(query.Start, 10) if err != nil { return nil, err @@ -100,10 +103,11 @@ func Search(query core.Query) ([]core.SearchResult, error) { } logrus.Debugf("Yandex URL built: %s", googleURL) - res, err := yandexRequest(googleURL, query) + res, err := yandexRequest(ctx, googleURL, query) if err != nil { return nil, err } + defer res.Body.Close() logrus.Debugf("Yandex Raw response: code=%d", res.StatusCode) results, err := yandexResultParser(res) diff --git a/yandex/search_test.go b/yandex/search_test.go index 23abcaf..38907e8 100644 --- a/yandex/search_test.go +++ b/yandex/search_test.go @@ -4,6 +4,7 @@ package yandex import ( + "context" "testing" "github.com/karust/openserp/core" @@ -18,7 +19,7 @@ func TestSearchYandex(t *testing.T) { yand := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "HEY", Limit: 10} - results, err := yand.Search(query) + results, err := yand.Search(context.Background(), query) ithelper.HandleError(t, "yandex web search", err) if len(results) == 0 { @@ -33,7 +34,7 @@ func TestImageYandex(t *testing.T) { yand := New(*browser, core.SearchEngineOptions{}) query := core.Query{Text: "furry tiger", Limit: 30} - results, err := yand.SearchImage(query) + results, err := yand.SearchImage(context.Background(), query) ithelper.HandleError(t, "yandex image search", err) if len(results) == 0 {