zlei6/openserp
1
0
Fork 0
mirror of https://github.com/karust/openserp.git synced 2026-08-05 16:53:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: propagate request context across search, retry, limiter, and browser navigation

Browse Source
This commit is contained in:
Rustem Kamalov
2026-04-20 05:53:13 +03:00
parent 0b7fbd3a45
commit e89b1a3a50
27 changed files with 424 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -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

11
baidu/search.go
View File

@@ -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
}

5
baidu/search_integration_test.go
View File

@@ -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 {

12
baidu/search_raw.go
View File

@@ -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)

30
bing/search.go
View File

@@ -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")

5
bing/search_integration_test.go
View File

@@ -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 {

2
cmd/root.go
View File

@@ -14,7 +14,7 @@ import (
)
const (
version = "0.6.3"
version = "0.6.4"
defaultConfigFilename = "config"
envPrefix = "OPENSERP"
)

10
cmd/search.go
View File

@@ -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")

19
cmd/serve.go
View File

@@ -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 {

54
core/browser.go
View File

@@ -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
}

5
core/browser_test.go
View File

@@ -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)
}

38
core/context.go Normal file
View File

@@ -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)
}

5
core/proxy_integration_test.go
View File

@@ -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
}

63
core/resilient.go
View File

@@ -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
}

75
core/resilient_context_test.go Normal file
View File

@@ -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")
}
}

35
core/retry.go
View File

@@ -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)
}

58
core/retry_test.go
View File

@@ -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)
}
}
}

17
core/server.go
View File

@@ -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 {

51
core/server_test.go
View File

@@ -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")
}

19
duckduckgo/search.go
View File

@@ -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

5
duckduckgo/search_integration_test.go
View File

@@ -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 {

15
google/search.go
View File

@@ -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
}

5
google/search_integration_test.go
View File

@@ -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 {

12
google/search_raw.go
View File

@@ -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)

15
yandex/search.go
View File

@@ -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
}

12
yandex/search_raw.go
View File

@@ -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)

5
yandex/search_test.go
View File

@@ -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 {