Stream image sizing, share client, honor context, cap concurrency

getImageSize buffered every image fully into memory via io.ReadAll, built
a fresh http.Client per image, ignored the request context, and spawned
one unbounded goroutine per <img> on the page.

It now streams to io.Discard behind a 10MB io.LimitReader (reading the body
rather than trusting Content-Length, so a lying header can't inflate an
image's rank), treats non-2xx responses as size 0 so an error page can't be
ranked as the lead image, shares a single lazily-built client, and
propagates the caller's context. extractPics caps concurrency at 8 probes
per page and bounds the whole probe phase with an overall budget (per-image
timeout kept below it) so a page full of slow image URLs can't hold the
handler past the write timeout; semaphore acquisition is context-aware.
Lead-image selection breaks size ties by URL so it's deterministic rather
than scheduling-dependent.

Author: paskal

extractor/pics.go

@@ -1,6 +1,7 @@
 package extractor
 
 import (
+	"context"
 	"io"
 	"net/http"
 	"sort"
@@ -11,21 +12,49 @@ import (
 	log "github.com/go-pkgz/lgr"
 )
 
-func (f *UReadability) extractPics(iselect *goquery.Selection, url string) (mainImage string, allImages []string, ok bool) {
-	images := make(map[int]string)
+const (
+	imageFetchTimeout       = 15 * time.Second // per-image probe timeout, kept below imageProbeBudget
+	imageProbeBudget        = 30 * time.Second // overall budget for probing every image on a page
+	maxImageBytes           = 10 << 20         // cap when streaming to measure, avoids buffering huge files
+	maxConcurrentImageFetch = 8                // limit parallel image probes per page
+)
 
-	type imgInfo struct {
-		url  string
-		size int
-	}
-	var resCh = make(chan imgInfo)
+// imageClient returns a lazily-built HTTP client for image size probes, shared across all fetches
+// instead of building a fresh client per image.
+func (f *UReadability) imageClient() *http.Client {
+	f.imgClientOnce.Do(func() {
+		f.imgClient = &http.Client{Timeout: imageFetchTimeout}
+	})
+	return f.imgClient
+}
+
+type imgInfo struct {
+	url  string
+	size int
+}
+
+func (f *UReadability) extractPics(ctx context.Context, iselect *goquery.Selection, url string) (mainImage string, allImages []string, ok bool) {
+	// bound total image probing so a page full of slow image URLs can't hold the handler open for
+	// ceil(n/maxConcurrentImageFetch) * imageFetchTimeout and blow past the server write timeout.
+	ctx, cancel := context.WithTimeout(ctx, imageProbeBudget)
+	defer cancel()
+
+	resCh := make(chan imgInfo)
 	var wg sync.WaitGroup
+	sem := make(chan struct{}, maxConcurrentImageFetch)
 
 	iselect.Each(func(_ int, s *goquery.Selection) {
-		if im, ok := s.Attr("src"); ok {
+		if im, exists := s.Attr("src"); exists {
 			wg.Go(func() {
-				size := f.getImageSize(im)
-				resCh <- imgInfo{url: im, size: size}
+				// acquire a slot, but give up if the overall budget is already spent
+				select {
+				case sem <- struct{}{}:
+					defer func() { <-sem }()
+				case <-ctx.Done():
+					resCh <- imgInfo{url: im, size: 0}
+					return
+				}
+				resCh <- imgInfo{url: im, size: f.getImageSize(ctx, im)}
 			})
 		}
 	})
@@ -35,37 +64,41 @@ func (f *UReadability) extractPics(iselect *goquery.Selection, url string) (main
 		close(resCh)
 	}()
 
+	var results []imgInfo
 	for r := range resCh {
-		images[r.size] = r.url
+		results = append(results, r)
 		allImages = append(allImages, r.url)
 	}
 	sort.Strings(allImages)
-	if len(images) == 0 {
+	if len(results) == 0 {
 		return "", nil, false
 	}
 
-	// get the biggest picture
-	keys := make([]int, 0, len(images))
-	for k := range images {
-		keys = append(keys, k)
+	// pick the largest image; break ties by URL so lead-image selection is deterministic rather
+	// than dependent on goroutine scheduling / map iteration order.
+	best := results[0]
+	for _, r := range results[1:] {
+		if r.size > best.size || (r.size == best.size && r.url < best.url) {
+			best = r
+		}
 	}
-	sort.Sort(sort.Reverse(sort.IntSlice(keys)))
-	mainImage = images[keys[0]]
-	log.Printf("[DEBUG] total images from %s = %d, main=%s (%d)", url, len(images), mainImage, keys[0])
-	return mainImage, allImages, true
+	log.Printf("[DEBUG] total images from %s = %d, main=%s (%d)", url, len(results), best.url, best.size)
+	return best.url, allImages, true
 }
 
-// getImageSize loads image to get size
-func (f *UReadability) getImageSize(url string) (size int) {
-	httpClient := &http.Client{Timeout: 30 * time.Second}
-	req, err := http.NewRequest("GET", url, http.NoBody)
+// getImageSize measures an image's byte size by streaming it to io.Discard behind a maxImageBytes
+// limit, without buffering it whole and honoring the caller's context. The body is actually read
+// (rather than trusting Content-Length) so a broken endpoint or a lying header can't inflate an
+// image's rank; an unreadable body sizes as 0, matching the previous behavior.
+func (f *UReadability) getImageSize(ctx context.Context, url string) (size int) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody)
 	if err != nil {
 		log.Printf("[WARN] can't create request to get pic from %s", url)
 		return 0
 	}
 	req.Close = true
 	req.Header.Set("User-Agent", userAgent)
-	resp, err := httpClient.Do(req)
+	resp, err := f.imageClient().Do(req)
 	if err != nil {
 		log.Printf("[WARN] can't get %s, error=%v", url, err)
 		return 0
@@ -76,10 +109,16 @@ func (f *UReadability) getImageSize(url string) (size int) {
 		}
 	}()
 
-	data, err := io.ReadAll(resp.Body)
+	// treat non-2xx as size 0 so an error page (404/500 HTML) can't be ranked as the lead image
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		log.Printf("[DEBUG] non-2xx (%d) for image %s, sizing as 0", resp.StatusCode, url)
+		return 0
+	}
+
+	n, err := io.Copy(io.Discard, io.LimitReader(resp.Body, maxImageBytes))
 	if err != nil {
 		log.Printf("[WARN] failed to get %s, err=%v", url, err)
 		return 0
 	}
-	return len(data)
+	return int(n)
 }

extractor/pics_test.go

@@ -49,7 +49,7 @@ func TestExtractPicsDirectly(t *testing.T) {
 		d, err := goquery.NewDocumentFromReader(strings.NewReader(data))
 		require.NoError(t, err)
 		sel := d.Find("img")
-		im, allImages, ok := lr.extractPics(sel, "url")
+		im, allImages, ok := lr.extractPics(context.Background(), sel, "url")
 		assert.True(t, ok)
 		assert.Len(t, allImages, 1)
 		assert.Equal(t, "https://cdn1.tnwcdn.com/wp-content/blogs.dir/1/files/2016/01/page-source.jpg", im)
@@ -60,7 +60,7 @@ func TestExtractPicsDirectly(t *testing.T) {
 		d, err := goquery.NewDocumentFromReader(strings.NewReader(data))
 		require.NoError(t, err)
 		sel := d.Find("img")
-		im, allImages, ok := lr.extractPics(sel, "url")
+		im, allImages, ok := lr.extractPics(context.Background(), sel, "url")
 		assert.False(t, ok)
 		assert.Empty(t, allImages)
 		assert.Empty(t, im)
@@ -71,12 +71,52 @@ func TestExtractPicsDirectly(t *testing.T) {
 		d, err := goquery.NewDocumentFromReader(strings.NewReader(data))
 		require.NoError(t, err)
 		sel := d.Find("img")
-		im, allImages, ok := lr.extractPics(sel, "url")
+		im, allImages, ok := lr.extractPics(context.Background(), sel, "url")
 		assert.True(t, ok)
 		assert.Len(t, allImages, 1)
 		assert.Equal(t, "http://bad_url", im)
 	})
 
+	t.Run("cancelled context returns zero size", func(t *testing.T) {
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			_, _ = w.Write([]byte("some-image-bytes"))
+		}))
+		defer ts.Close()
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		assert.Equal(t, 0, lr.getImageSize(ctx, ts.URL))
+	})
+
+	t.Run("measures image size by streamed bytes", func(t *testing.T) {
+		payload := strings.Repeat("x", 1234)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			_, _ = w.Write([]byte(payload))
+		}))
+		defer ts.Close()
+		assert.Equal(t, len(payload), lr.getImageSize(context.Background(), ts.URL))
+	})
+
+	t.Run("streams and counts bytes when length unknown", func(t *testing.T) {
+		payload := strings.Repeat("x", 1234)
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			fl, isFlusher := w.(http.Flusher)
+			require.True(t, isFlusher)
+			_, _ = w.Write([]byte(payload[:600]))
+			fl.Flush() // forces chunked transfer, so Content-Length is unknown on the client
+			_, _ = w.Write([]byte(payload[600:]))
+		}))
+		defer ts.Close()
+		assert.Equal(t, len(payload), lr.getImageSize(context.Background(), ts.URL))
+	})
+
+	t.Run("non-2xx response sizes as zero", func(t *testing.T) {
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			http.Error(w, "not found", http.StatusNotFound) // error page body must not be ranked as an image
+		}))
+		defer ts.Close()
+		assert.Equal(t, 0, lr.getImageSize(context.Background(), ts.URL))
+	})
+
 	t.Run("bad body of the image", func(t *testing.T) {
 		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Content-Length", "1") // error on reading body
@@ -86,7 +126,7 @@ func TestExtractPicsDirectly(t *testing.T) {
 		d, err := goquery.NewDocumentFromReader(strings.NewReader(data))
 		require.NoError(t, err)
 		sel := d.Find("img")
-		im, allImages, ok := lr.extractPics(sel, "url")
+		im, allImages, ok := lr.extractPics(context.Background(), sel, "url")
 		assert.True(t, ok)
 		assert.Len(t, allImages, 1)
 		assert.Equal(t, ts.URL, im)

extractor/readability.go

@@ -4,6 +4,7 @@ package extractor
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"net/url"
 	"regexp"
 	"strings"
@@ -40,6 +41,9 @@ type UReadability struct {
 
 	defaultRetrieverOnce sync.Once
 	defaultRetriever     Retriever
+
+	imgClientOnce sync.Once
+	imgClient     *http.Client
 }
 
 // retriever returns the configured default Retriever, creating a cached HTTPRetriever if nil
@@ -153,7 +157,7 @@ func (f *UReadability) extractWithRules(ctx context.Context, reqURL string, rule
 		log.Printf("[WARN] failed to create document from reader, error=%v", err)
 		return nil, err
 	}
-	if im, allImages, ok := f.extractPics(darticle.Find("img"), reqURL); ok {
+	if im, allImages, ok := f.extractPics(ctx, darticle.Find("img"), reqURL); ok {
 		rb.Image = im
 		rb.AllImages = allImages
 	}