From bd4c6d21dd00a1e52c3dfd3c55649dfb931dfc04 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 16 Apr 2020 23:24:33 +0300
Subject: [PATCH] lib/promscrape: retry target scraping when the target closes
 previously established keep-alive connection to it

This should fix the following error:

the server closed connection before returning the first response byte. Make sure the server returns 'Connection: close' response header before closing the connection
---
 lib/promscrape/client.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/lib/promscrape/client.go b/lib/promscrape/client.go
index 01334b02f4..dcb265ee55 100644
--- a/lib/promscrape/client.go
+++ b/lib/promscrape/client.go
@@ -79,8 +79,8 @@ func (c *client) ReadData(dst []byte) ([]byte, error) {
 		req.Header.Set("Authorization", c.authHeader)
 	}
 	resp := fasthttp.AcquireResponse()
-	// There is no need in calling DoTimeout, since the timeout is already set in c.hc.ReadTimeout.
-	err := c.hc.Do(req, resp)
+	err := doRequestWithPossibleRetry(c.hc, req, resp)
+
 	fasthttp.ReleaseRequest(req)
 	if err != nil {
 		fasthttp.ReleaseResponse(resp)
@@ -120,3 +120,16 @@ var (
 	scrapesGunzipped    = metrics.NewCounter(`vm_promscrape_scrapes_gunziped_total`)
 	scrapesGunzipFailed = metrics.NewCounter(`vm_promscrape_scrapes_gunzip_failed_total`)
 )
+
+func doRequestWithPossibleRetry(hc *fasthttp.HostClient, req *fasthttp.Request, resp *fasthttp.Response) error {
+	// There is no need in calling DoTimeout, since the timeout must be already set in hc.ReadTimeout.
+	err := hc.Do(req, resp)
+	if err == nil {
+		return nil
+	}
+	if err != fasthttp.ErrConnectionClosed {
+		return err
+	}
+	// Retry request if the server closed the keep-alive connection during the first attempt.
+	return hc.Do(req, resp)
+}