Bypass request delay when request is cancelled

extensions/random_user_agent.go

@@ -264,10 +264,10 @@ func genChromeUA() string {
 //	-> "User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.72 Safari/537.36 Edg/90.0.818.39"
 func genEdgeUA() string {
 	version := edgeVersions[rand.Intn(len(edgeVersions))]
-	chrome_version := strings.Split(version, ",")[0]
-	edge_version := strings.Split(version, ",")[1]
+	chromeVersion := strings.Split(version, ",")[0]
+	edgeVersion := strings.Split(version, ",")[1]
 	os := osStrings[rand.Intn(len(osStrings))]
-	return fmt.Sprintf("Mozilla/5.0 (%s) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36 Edg/%s", os, chrome_version, edge_version)
+	return fmt.Sprintf("Mozilla/5.0 (%s) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%s Safari/537.36 Edg/%s", os, chromeVersion, edgeVersion)
 }
 
 // Generates Opera Browser User-Agent (Desktop)

http_backend.go

@@ -18,6 +18,7 @@ import (
 	"crypto/sha1"
 	"encoding/gob"
 	"encoding/hex"
+	"errors"
 	"io"
 	"io/ioutil"
 	"math/rand"
@@ -170,7 +171,12 @@ func (h *httpBackend) Cache(request *http.Request, bodySize int, checkHeadersFun
 func (h *httpBackend) Do(request *http.Request, bodySize int, checkHeadersFunc checkHeadersFunc) (*Response, error) {
 	r := h.GetMatchingRule(request.URL.Host)
 	if r != nil {
-		r.waitChan <- true
+		select {
+		case r.waitChan <- true:
+		case <-request.Context().Done():
+			return nil, errors.New("context canceled; early return")
+		}
+
 		defer func(r *LimitRule) {
 			randomDelay := time.Duration(0)
 			if r.RandomDelay != 0 {

http_backend_test.go

@@ -0,0 +1,95 @@
+package colly
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"net/http"
+	"net/http/cookiejar"
+	"net/http/httptest"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestHTTPBackendDoCancelation(t *testing.T) {
+	// hardcoded fixed to ensure that p, n, c are the same making the test more reliable
+	rand.Seed(2927496558871806)
+
+	// rand up to 10 to not extend the test duration too much
+	p := 1 + rand.Intn(5)        // p: parallel requests
+	n := p + p*rand.Intn(10)     // n: after n, cancel will be called; ensure 1 calls per worker + rand
+	c := n + p*2 + rand.Intn(10) // c: total number of calls; ensure 2 calls per worker after cancel is called + rand
+
+	ts := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(rw, "OK")
+	}))
+	defer ts.Close()
+
+	checkHeadersFunc := func(req *http.Request, statusCode int, header http.Header) bool { return true }
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	backend := &httpBackend{}
+	jar, _ := cookiejar.New(nil)
+	backend.Init(jar)
+	limit := &LimitRule{
+		DomainRegexp: ".*",
+		Parallelism:  p,
+		Delay:        5 * time.Millisecond,
+	}
+	backend.Limit(limit)
+
+	var wg sync.WaitGroup
+	wg.Add(c)
+
+	out := make(chan error)
+
+	for i := 0; i < c; i++ {
+		go func(i int) {
+			defer wg.Done()
+			req, _ := http.NewRequest("GET", ts.URL+"/"+strconv.Itoa(i), nil)
+			req = req.WithContext(ctx)
+
+			_, err := backend.Do(req, 0, checkHeadersFunc)
+			out <- err
+		}(i)
+	}
+
+	go func() {
+		wg.Wait()
+		close(out)
+	}()
+
+	i := 0
+	nonEarlyCount := 0
+	for err := range out {
+		i++
+		if i == n {
+			cancel()
+		}
+
+		if i <= n {
+			if err != nil {
+				t.Errorf("no error was expected for the first %d responses; error: %q", n, err)
+			}
+		} else {
+			errStr := ""
+			if err != nil {
+				errStr = err.Error()
+			}
+
+			// non early returns are allowed up to the number of maximum allowed concurrent requests;
+			// bacause those requests could be already running when cancel was called
+			if !strings.Contains(errStr, "early return") {
+				if nonEarlyCount > p {
+					t.Errorf("count of non early return is above the number of maximum allowed concurrent requests; p: %d; n: %d; c: %d", p, n, c)
+				}
+				nonEarlyCount++
+			}
+		}
+	}
+}