Repository: gocolly/colly

Issue:
need a Callback func after LimitRule but before h.Client.Do
<!--
Love colly? Please consider supporting our collective:
👉  https://opencollective.com/colly/donate
-->
When handling concurrent requests, multiple identical requests may wait for the LimitRule. Even if one request succeeds at the moment, other requests will still be processed.

We need a callback method that will be invoked immediately after the waiting of LimitRule is completed, to determine whether the request needs to be terminated, similar to checkHeadersFunc.


当并发请求的时候 会出现相同请求等待LimitRule,如果某一个请求此刻成功后,其他请求依然会被请求
需要一个回调方法,在LimitRule等待完毕后立马进行调用,判断是否需要终止请求,,类似checkHeadersFunc
 

Discussion:

Solution patch:
diff --git a/colly.go b/colly.go
index 833ca4f..af045f7 100644
--- a/colly.go
+++ b/colly.go
@@ -128,6 +128,7 @@ type Collector struct {
 	requestCallbacks         []RequestCallback
 	responseCallbacks        []ResponseCallback
 	responseHeadersCallbacks []ResponseHeadersCallback
+	requestHeadersCallbacks  []RequestCallback
 	errorCallbacks           []ErrorCallback
 	scrapedCallbacks         []ScrapedCallback
 	requestCount             uint32
@@ -235,6 +236,8 @@ var (
 	ErrEmptyProxyURL = errors.New("Proxy URL list is empty")
 	// ErrAbortedAfterHeaders is the error returned when OnResponseHeaders aborts the transfer.
 	ErrAbortedAfterHeaders = errors.New("Aborted after receiving response headers")
+	// ErrAbortedBeforeRequest is the error returned when OnResponseHeaders aborts the transfer.
+	ErrAbortedBeforeRequest = errors.New("Aborted before Do Request")
 	// ErrQueueFull is the error returned when the queue is full
 	ErrQueueFull = errors.New("Queue MaxSize reached")
 	// ErrMaxRequests is the error returned when exceeding max requests
@@ -714,7 +717,7 @@ func (c *Collector) fetch(u, method string, depth int, requestData io.Reader, ct
 		req = hTrace.WithTrace(req)
 	}
 	origURL := req.URL
-	checkHeadersFunc := func(req *http.Request, statusCode int, headers http.Header) bool {
+	checkResponseHeadersFunc := func(req *http.Request, statusCode int, headers http.Header) bool {
 		if req.URL != origURL {
 			request.URL = req.URL
 			request.Headers = &req.Header
@@ -722,7 +725,11 @@ func (c *Collector) fetch(u, method string, depth int, requestData io.Reader, ct
 		c.handleOnResponseHeaders(&Response{Ctx: ctx, Request: request, StatusCode: statusCode, Headers: &headers})
 		return !request.abort
 	}
-	response, err := c.backend.Cache(req, c.MaxBodySize, checkHeadersFunc, c.CacheDir, c.CacheExpiration)
+	checkRequestHeadersFunc := func(req *http.Request) bool {
+		c.handleOnRequestHeaders(request)
+		return !request.abort
+	}
+	response, err := c.backend.Cache(req, c.MaxBodySize, checkRequestHeadersFunc, checkResponseHeadersFunc, c.CacheDir, c.CacheExpiration)
 	if proxyURL, ok := req.Context().Value(ProxyURLKey).(string); ok {
 		request.ProxyURL = proxyURL
 	}
@@ -944,6 +951,14 @@ func (c *Collector) OnResponseHeaders(f ResponseHeadersCallback) {
 	c.lock.Unlock()
 }
 
+// OnRequestHeaders registers a function. Function will be executed on every
+// request made by the Collector before Request Do
+func (c *Collector) OnRequestHeaders(f RequestCallback) {
+	c.lock.Lock()
+	c.requestHeadersCallbacks = append(c.requestHeadersCallbacks, f)
+	c.lock.Unlock()
+}
+
 // OnResponse registers a function. Function will be executed on every response
 func (c *Collector) OnResponse(f ResponseCallback) {
 	c.lock.Lock()
@@ -1149,6 +1164,16 @@ func (c *Collector) handleOnResponseHeaders(r *Response) {
 		f(r)
 	}
 }
+func (c *Collector) handleOnRequestHeaders(r *Request) {
+	if c.debugger != nil {
+		c.debugger.Event(createEvent("requestHeaders", r.ID, c.ID, map[string]string{
+			"url": r.URL.String(),
+		}))
+	}
+	for _, f := range c.requestHeadersCallbacks {
+		f(r)
+	}
+}
 
 func (c *Collector) handleOnHTML(resp *Response) error {
 	c.lock.RLock()
diff --git a/http_backend.go b/http_backend.go
index 3d178f5..9404e7d 100644
--- a/http_backend.go
+++ b/http_backend.go
@@ -39,7 +39,8 @@ type httpBackend struct {
 	lock       *sync.RWMutex
 }
 
-type checkHeadersFunc func(req *http.Request, statusCode int, header http.Header) bool
+type checkResponseHeadersFunc func(req *http.Request, statusCode int, header http.Header) bool
+type checkRequestHeadersFunc func(req *http.Request) bool
 
 // LimitRule provides connection restrictions for domains.
 // Both DomainRegexp and DomainGlob can be used to specify
@@ -128,9 +129,9 @@ func (h *httpBackend) GetMatchingRule(domain string) *LimitRule {
 	return nil
 }
 
-func (h *httpBackend) Cache(request *http.Request, bodySize int, checkHeadersFunc checkHeadersFunc, cacheDir string, cacheExpiration time.Duration) (*Response, error) {
+func (h *httpBackend) Cache(request *http.Request, bodySize int, checkRequestHeadersFunc checkRequestHeadersFunc, checkResponseHeadersFunc checkResponseHeadersFunc, cacheDir string, cacheExpiration time.Duration) (*Response, error) {
 	if cacheDir == "" || request.Method != "GET" || request.Header.Get("Cache-Control") == "no-cache" {
-		return h.Do(request, bodySize, checkHeadersFunc)
+		return h.Do(request, bodySize, checkRequestHeadersFunc, checkResponseHeadersFunc)
 	}
 	sum := sha1.Sum([]byte(request.URL.String()))
 	hash := hex.EncodeToString(sum[:])
@@ -147,12 +148,12 @@ func (h *httpBackend) Cache(request *http.Request, bodySize int, checkHeadersFun
 		resp := new(Response)
 		err := gob.NewDecoder(file).Decode(resp)
 		file.Close()
-		checkHeadersFunc(request, resp.StatusCode, *resp.Headers)
+		checkResponseHeadersFunc(request, resp.StatusCode, *resp.Headers)
 		if resp.StatusCode < 500 {
 			return resp, err
 		}
 	}
-	resp, err := h.Do(request, bodySize, checkHeadersFunc)
+	resp, err := h.Do(request, bodySize, checkRequestHeadersFunc, checkResponseHeadersFunc)
 	if err != nil || resp.StatusCode >= 500 {
 		return resp, err
 	}
@@ -173,7 +174,7 @@ func (h *httpBackend) Cache(request *http.Request, bodySize int, checkHeadersFun
 	return resp, os.Rename(filename+"~", filename)
 }
 
-func (h *httpBackend) Do(request *http.Request, bodySize int, checkHeadersFunc checkHeadersFunc) (*Response, error) {
+func (h *httpBackend) Do(request *http.Request, bodySize int, checkRequestHeadersFunc checkRequestHeadersFunc, checkResponseHeadersFunc checkResponseHeadersFunc) (*Response, error) {
 	r := h.GetMatchingRule(request.URL.Host)
 	if r != nil {
 		r.waitChan <- true
@@ -186,7 +187,9 @@ func (h *httpBackend) Do(request *http.Request, bodySize int, checkHeadersFunc c
 			<-r.waitChan
 		}(r)
 	}
-
+	if !checkRequestHeadersFunc(request) {
+		return nil, ErrAbortedBeforeRequest
+	}
 	res, err := h.Client.Do(request)
 	if err != nil {
 		return nil, err
@@ -197,7 +200,7 @@ func (h *httpBackend) Do(request *http.Request, bodySize int, checkHeadersFunc c
 	if res.Request != nil {
 		finalRequest = res.Request
 	}
-	if !checkHeadersFunc(finalRequest, res.StatusCode, res.Header) {
+	if !checkResponseHeadersFunc(finalRequest, res.StatusCode, res.Header) {
 		// closing res.Body (see defer above) without reading it aborts
 		// the download
 		return nil, ErrAbortedAfterHeaders



Test patch:
diff --git a/colly_test.go b/colly_test.go
index 5229388..bae4e8a 100644
--- a/colly_test.go
+++ b/colly_test.go
@@ -1884,3 +1884,23 @@ func TestCollectorPostRetryUnseekable(t *testing.T) {
 		t.Error("OnResponse Retry was called but BodyUnseekable")
 	}
 }
+
+func TestCheckRequestHeadersFunc(t *testing.T) {
+	ts := newTestServer()
+	defer ts.Close()
+	try := false
+
+	c := NewCollector()
+
+	c.OnRequestHeaders(func(r *Request) {
+		try = true
+		r.Abort()
+	})
+	c.OnScraped(func(r *Response) {
+		try = false
+	})
+	c.Visit(ts.URL)
+	if try == false {
+		t.Error("TestCheckRequestHeadersFunc failed")
+	}
+}