app/vmauth: simplify the code after 4a1d29126c

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4242

app/vmauth/main.go

@@ -158,8 +158,9 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 		up, headers = ui.DefaultURL, ui.Headers
 		isDefault = true
 	}
-	rtb := &readTrackingBody{ioc: r.Body}
+	r.Body = &readTrackingBody{
-	r.Body = rtb
+		r: r.Body,
+	}
 
 	maxAttempts := up.getBackendsCount()
 	for i := 0; i < maxAttempts; i++ {
@@ -180,7 +181,6 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 		}
 		bu.setBroken()
 	}
-	rtb.mustClose()
 	err := &httpserver.ErrorWithStatusCode{
 		Err:        fmt.Errorf("all the backends for the user %q are unavailable", ui.name()),
 		StatusCode: http.StatusServiceUnavailable,
@@ -195,16 +195,13 @@ func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url
 	for _, h := range headers {
 		req.Header.Set(h.Name, h.Value)
 	}
-	rtb := req.Body.(*readTrackingBody)
 	transportOnce.Do(transportInit)
 	res, err := transport.RoundTrip(req)
 	if err != nil {
-		remoteAddr := httpserver.GetQuotedRemoteAddr(r)
+		rtb := req.Body.(*readTrackingBody)
-		// NOTE: do not use httpserver.GetRequestURI
+		if rtb.readStarted {
-		// it explicitly reads request body and fail retries
+			// Request body has been already read, so it is impossible to retry the request.
-		if (r.Method == http.MethodPost || r.Method == http.MethodPut) && !rtb.couldReuseBody() {
+			// Return the error to the client then.
-			// It is impossible to retry POST and PUT requests,
-			// since we already proxied the request body to the backend.
 			err = &httpserver.ErrorWithStatusCode{
 				Err:        fmt.Errorf("cannot proxy the request to %q: %w", targetURL, err),
 				StatusCode: http.StatusServiceUnavailable,
@@ -212,6 +209,10 @@ func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url
 			httpserver.Errorf(w, r, "%s", err)
 			return true
 		}
+		// Retry the request if its body wasn't read yet. This usually means that the backend isn't reachable.
+		remoteAddr := httpserver.GetQuotedRemoteAddr(r)
+		// NOTE: do not use httpserver.GetRequestURI
+		// it explicitly reads request body and fails retries.
 		logger.Warnf("remoteAddr: %s; requestURI: %s; error when proxying the request to %q: %s", remoteAddr, req.URL, targetURL, err)
 		return false
 	}
@@ -223,7 +224,6 @@ func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url
 	copyBuf.B = bytesutil.ResizeNoCopyNoOverallocate(copyBuf.B, 16*1024)
 	_, err = io.CopyBuffer(w, res.Body, copyBuf.B)
 	copyBufPool.Put(copyBuf)
-	rtb.mustClose()
 	if err != nil && !netutil.IsTrivialNetworkError(err) {
 		remoteAddr := httpserver.GetQuotedRemoteAddr(r)
 		requestURI := httpserver.GetRequestURI(r)
@@ -357,41 +357,26 @@ func handleConcurrencyLimitError(w http.ResponseWriter, r *http.Request, err err
 }
 
 type readTrackingBody struct {
-	ioc       io.ReadCloser
+	r           io.ReadCloser
-	wasClosed bool
+	readStarted bool
-	wasRead   bool
 }
 
 // Read implements io.Reader interface
 // tracks body reading requests
 func (rtb *readTrackingBody) Read(p []byte) (int, error) {
-	if rtb.wasClosed {
+	if len(p) > 0 {
-		return 0, io.EOF
+		rtb.readStarted = true
 	}
-	rtb.wasRead = true
+	return rtb.r.Read(p)
-	return rtb.ioc.Read(p)
 }
 
-// Close implements interface.
+// Close implements io.Closer interface.
-// Closes body only if Read call was performed.
-// http.Roundtrip performs body.Close call even without any Read calls
-// so this hack allows us to reuse request body
 func (rtb *readTrackingBody) Close() error {
-	if rtb.wasRead {
+	// Close rtb.r only if at least a single Read call was performed.
-		err := rtb.ioc.Close()
+	// http.Roundtrip performs body.Close call even without any Read calls
-		rtb.wasClosed = true
+	// so this hack allows us to reuse request body
-		return err
+	if rtb.readStarted {
+		return rtb.r.Close()
 	}
 	return nil
 }
-
-// body could be reused only if read operation wasn't called
-func (rtb *readTrackingBody) couldReuseBody() bool {
-	return !rtb.wasRead
-}
-
-// mustClose explicitly closes body
-func (rtb *readTrackingBody) mustClose() {
-	rtb.wasClosed = true
-	rtb.ioc.Close()
-}

docs/CHANGELOG.md

@@ -46,7 +46,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 * FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): add ability to filter incoming requests by IP. See [these docs](https://docs.victoriametrics.com/vmauth.html#ip-filters) and [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3491).
 * FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): add ability to proxy requests to the specified backends for unauthorized users. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4083).
 * FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): add ability to specify default route for unmatched requests. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4084).
-* FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): add ability to retry requests for common tcp dial error for POST and PUT requests. Previously load-balancing algorythm didn't retry such errors.
+* FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): retry `POST` requests on the remaining backends if the currently selected backend isn't reachable. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4242).
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): add ability to compare the data for the previous day with the data for the current day at [Cardinality Explorer](https://docs.victoriametrics.com/#cardinality-explorer). See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3967).
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): display histograms as heatmaps in [Metrics explorer](https://docs.victoriametrics.com/#metrics-explorer). See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4111).
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): add `WITH template` playground. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3811).