vmagent: maintenance mode support

app/vmagent/main.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"embed"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"io"
 	"net/http"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -549,6 +551,33 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool {
 		procutil.SelfSIGHUP()
 		w.WriteHeader(http.StatusOK)
 		return true
+	case "/remotewrite/maintenance":
+		if !remotewrite.CheckMaintenanceAuthKey(w, r) {
+			return true
+		}
+		remoteWriteMaintenanceRequests.Inc()
+		if v := r.FormValue("enable"); v != "" {
+			enable, err := strconv.ParseBool(v)
+			if err != nil {
+				httpserver.Errorf(w, r, "cannot parse `enable` query arg: %s", err)
+				return true
+			}
+			rwURL := r.FormValue("url")
+			if rwURL == "" {
+				httpserver.Errorf(w, r, "missing `url` query arg; it must match the `url` label of vmagent_remotewrite_* metrics for the target -remoteWrite.url, "+
+					"or be set to `*` to target all the configured -remoteWrite.url destinations")
+				return true
+			}
+			if matched := remotewrite.SetMaintenanceMode(rwURL, enable); matched == 0 {
+				httpserver.Errorf(w, r, "no -remoteWrite.url destinations match `url=%s`", rwURL)
+				return true
+			}
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"maintenance": remotewrite.GetMaintenanceMode(),
+		})
+		return true
 	case "/ready":
 		if rdy := promscrape.PendingScrapeConfigs.Load(); rdy > 0 {
 			errMsg := fmt.Sprintf("waiting for scrapes to init, left: %d", rdy)
@@ -833,6 +862,8 @@ var (
 	remoteWriteStatusURLRelabelConfigRequests = metrics.NewCounter(`vmagent_http_requests_total{path="/api/v1/status/remotewrite-url-relabel-config"}`)
 
 	promscrapeConfigReloadRequests = metrics.NewCounter(`vmagent_http_requests_total{path="/-/reload"}`)
+
+	remoteWriteMaintenanceRequests = metrics.NewCounter(`vmagent_http_requests_total{path="/remotewrite/maintenance"}`)
 )
 
 func usage() {

app/vmagent/remotewrite/client.go

@@ -97,6 +97,8 @@ type client struct {
 	useVMProto          atomic.Bool
 	canDowngradeVMProto atomic.Bool
 
+	maintenanceMode atomic.Bool
+
 	fq *persistentqueue.FastQueue
 	hc *http.Client
 
@@ -322,6 +324,12 @@ func (c *client) runWorker(readBlock func(dst []byte) ([]byte, bool)) {
 	var block []byte
 	ch := make(chan bool, 1)
 	for {
+		if c.maintenanceMode.Load() {
+			if !c.waitForMaintenanceModeOff() {
+				return
+			}
+		}
+
 		block, ok = readBlock(block[:0])
 		if !ok {
 			return
@@ -339,6 +347,14 @@ func (c *client) runWorker(readBlock func(dst []byte) ([]byte, bool)) {
 			}
 			// Return unsent block to the queue.
 			c.fq.MustWriteBlockIgnoreDisabledPQ(block)
+			select {
+			case <-c.stopCh:
+				// c must be stopped.
+			default:
+				// sendBlock returned false because maintenance mode is enabled, not because
+				// c is stopping. Keep the worker alive so it resumes once maintenance mode is disabled.
+				continue
+			}
 			return
 		case <-c.stopCh:
 			// c must be stopped. Wait up to 5 seconds for the in-flight request to complete.
@@ -363,6 +379,23 @@ func (c *client) runWorker(readBlock func(dst []byte) ([]byte, bool)) {
 	}
 }
 
+// waitForMaintenanceModeOff blocks while maintenance mode is enabled for c.
+//
+// It returns false only if c.stopCh is closed while waiting.
+func (c *client) waitForMaintenanceModeOff() bool {
+	t := time.NewTicker(time.Second)
+	defer t.Stop()
+
+	for c.maintenanceMode.Load() {
+		select {
+		case <-t.C:
+		case <-c.stopCh:
+			return false
+		}
+	}
+	return true
+}
+
 func (c *client) doRequest(url string, body []byte) (*http.Response, error) {
 	req, err := c.newRequest(url, body)
 	if err != nil {
@@ -421,7 +454,7 @@ func (c *client) newRequest(url string, body []byte) (*http.Request, error) {
 
 // sendBlockHTTP sends the given block to c.remoteWriteURL.
 //
-// The function returns false only if c.stopCh is closed.
+// The function returns false if c.stopCh is closed or if maintenance mode is enabled for c.
 // Otherwise, it tries sending the block to remote storage indefinitely.
 func (c *client) sendBlockHTTP(block []byte) bool {
 	c.rl.Register(len(block))
@@ -429,6 +462,10 @@ func (c *client) sendBlockHTTP(block []byte) bool {
 	retriesCount := 0
 
 again:
+	if c.maintenanceMode.Load() {
+		return false
+	}
+
 	startTime := time.Now()
 	resp, err := c.doRequest(c.remoteWriteURL, block)
 	c.requestDuration.UpdateDuration(startTime)
@@ -511,12 +548,7 @@ again:
 	// Handle response
 	body, err := io.ReadAll(resp.Body)
 	_ = resp.Body.Close()
-	if err != nil {
-		logger.Errorf("cannot read response body from %q during retry #%d: %s", c.sanitizedURL, retriesCount, err)
-	} else {
-		logger.Errorf("unexpected status code received after sending a block with size %d bytes to %q during retry #%d: %d; response body=%q; "+
-			"re-sending the block in %s", len(block), c.sanitizedURL, retriesCount, statusCode, body, bt.CurrentDelay())
-	}
+	logUnexpectedStatusCode(block, c.sanitizedURL, statusCode, retriesCount, bt.CurrentDelay(), retryAfterHeader > 0, body, err)
 	if !bt.Wait(c.stopCh) {
 		return false
 	}
@@ -552,6 +584,26 @@ func (c *client) drainInMemoryQueue(stopCtx context.Context, block []byte) {
 
 var remoteWriteRejectedLogger = logger.WithThrottler("remoteWriteRejected", 5*time.Second)
 var remoteWriteRetryLogger = logger.WithThrottler("remoteWriteRetry", 5*time.Second)
+var remoteWriteUnexpectedStatusLogger = logger.WithThrottler("remoteWriteUnexpectedStatus", 5*time.Second)
+
+func logUnexpectedStatusCode(block []byte, sanitizedURL string, statusCode, retriesCount int, retryDelay time.Duration, isExpectedBackoff bool, body []byte, bodyErr error) {
+	if bodyErr != nil {
+		remoteWriteUnexpectedStatusLogger.Errorf("cannot read response body from %q during retry #%d: %s", sanitizedURL, retriesCount, bodyErr)
+		return
+	}
+
+	msg := fmt.Sprintf("unexpected status code received after sending a block with size %d bytes to %q during retry #%d: %d; response body=%q; "+
+		"re-sending the block in %s", len(block), sanitizedURL, retriesCount, statusCode, body, retryDelay)
+
+	if isExpectedBackoff {
+		// The remote storage explicitly signaled backoff duration via the Retry-After header,
+		// so this isn't an anomaly worth an ERROR log.
+		remoteWriteUnexpectedStatusLogger.Warnf("%s", msg)
+		return
+	}
+
+	remoteWriteUnexpectedStatusLogger.Errorf("%s", msg)
+}
 
 // repackBlockFromZstdToSnappy repacks the given zstd-compressed block to snappy-compressed block.
 //

app/vmagent/remotewrite/remotewrite.go

@@ -113,6 +113,7 @@ var (
 		"Multiple label names should be separated by `^^`, e.g. \"job^^instance,ip\". "+
 		"Can be combined with -remoteWrite.mdx.enable to hide sensitive label values in VictoriaMetrics self-monitoring metrics. "+
 		"Please see https://docs.victoriametrics.com/victoriametrics/vmagent/#obfuscating-label-values")
+	maintenanceAuthKey = flagutil.NewPassword("remoteWrite.maintenanceAuthKey", "Auth key for /remotewrite/maintenance http endpoint. It must be passed via authKey query arg. It overrides -httpAuth.*")
 )
 
 var (
@@ -141,6 +142,42 @@ func MultitenancyEnabled() bool {
 	return *enableMultitenantHandlers
 }
 
+// SetMaintenanceMode enables or disables maintenance mode for the remote write destination(s)
+// identified by sanitizedURL, which must match the `url` label of the corresponding
+// vmagent_remotewrite_* metrics. Pass "*" to target all the configured -remoteWrite.url destinations.
+//
+// While a destination is in maintenance mode, vmagent doesn't attempt to send data to it at all -
+// it stops draining that destination's queue; buffering remains subject to the configured queue limits.
+//
+// It returns the number of destinations matched by sanitizedURL.
+func SetMaintenanceMode(sanitizedURL string, enabled bool) int {
+	matched := 0
+	for _, rwctx := range rwctxsGlobal {
+		if sanitizedURL != "*" && rwctx.c.sanitizedURL != sanitizedURL {
+			continue
+		}
+		rwctx.c.maintenanceMode.Store(enabled)
+		matched++
+	}
+	return matched
+}
+
+// GetMaintenanceMode returns the maintenance mode status for each configured -remoteWrite.url,
+// keyed by its sanitized URL.
+func GetMaintenanceMode() map[string]bool {
+	m := make(map[string]bool, len(rwctxsGlobal))
+	for _, rwctx := range rwctxsGlobal {
+		m[rwctx.c.sanitizedURL] = rwctx.c.maintenanceMode.Load()
+	}
+	return m
+}
+
+// CheckMaintenanceAuthKey verifies the authKey query arg for the /remotewrite/maintenance http endpoint
+// against -remoteWrite.maintenanceAuthKey. See httpserver.CheckAuthFlag for the semantics of the return value.
+func CheckMaintenanceAuthKey(w http.ResponseWriter, r *http.Request) bool {
+	return httpserver.CheckAuthFlag(w, r, maintenanceAuthKey)
+}
+
 // Contains the current relabelConfigs.
 var allRelabelConfigs atomic.Pointer[relabelConfigs]
 

lib/httpserver/httpserver.go

@@ -488,7 +488,8 @@ func isProtectedByAuthFlag(path string) bool {
 	return strings.HasSuffix(path, "/config") || strings.HasSuffix(path, "/reload") ||
 		strings.HasSuffix(path, "/resetRollupResultCache") || strings.HasSuffix(path, "/delSeries") || strings.HasSuffix(path, "/delete_series") ||
 		strings.HasSuffix(path, "/force_merge") || strings.HasSuffix(path, "/force_flush") || strings.HasSuffix(path, "/snapshot") ||
-		strings.HasPrefix(path, "/snapshot/") || strings.HasSuffix(path, "/admin/status/metric_names_stats/reset")
+		strings.HasPrefix(path, "/snapshot/") || strings.HasSuffix(path, "/admin/status/metric_names_stats/reset") ||
+		strings.HasSuffix(path, "/remotewrite/maintenance")
 }
 
 // CheckAuthFlag checks whether the given authKey is set and valid