Move remote write API to client_golang/exp

Makefile

@@ -14,12 +14,16 @@
 include .bingo/Variables.mk
 include Makefile.common
 
+.PHONY: deps
+deps:
+	$(GO) work sync
+	$(MAKE) common-deps
+
 .PHONY: test
-test: deps common-test
+test: deps common-test test-exp
 
 .PHONY: test-short
-test-short: deps common-test-short
-
+test-short: deps common-test-short test-exp-short
 # Overriding Makefile.common check_license target to add
 # dagger paths
 .PHONY: common-check_license
@@ -55,8 +59,17 @@ fmt: common-format $(GOIMPORTS)
 proto: ## Regenerate Go from remote write proto.
 proto: $(BUF)
 	@echo ">> regenerating Prometheus Remote Write proto"
-	@cd api/prometheus/v1/remote/genproto && $(BUF) generate
-	@cd api/prometheus/v1/remote && find genproto/ -type f -exec sed -i '' 's/protohelpers "github.com\/planetscale\/vtprotobuf\/protohelpers"/protohelpers "github.com\/prometheus\/client_golang\/internal\/github.com\/planetscale\/vtprotobuf\/protohelpers"/g' {} \;
+	@cd exp/api/remote/genproto && $(BUF) generate
+	@cd exp/api/remote && find genproto/ -type f -exec sed -i '' 's/protohelpers "github.com\/planetscale\/vtprotobuf\/protohelpers"/protohelpers "github.com\/prometheus\/client_golang\/exp\/internal\/github.com\/planetscale\/vtprotobuf\/protohelpers"/g' {} \;
 	# For some reasons buf generates this unused import, kill it manually for now and reformat.
-	@cd api/prometheus/v1/remote && find genproto/ -type f -exec sed -i '' 's/_ "github.com\/gogo\/protobuf\/gogoproto"//g' {} \;
-	@cd api/prometheus/v1/remote && go fmt ./genproto/...
+	@cd exp/api/remote && find genproto/ -type f -exec sed -i '' 's/_ "github.com\/gogo\/protobuf\/gogoproto"//g' {} \;
+	@cd exp/api/remote && go fmt ./genproto/...
+	$(MAKE) fmt
+
+.PHONY: test-exp
+test-exp:
+	cd exp && $(GOTEST) $(test-flags) $(GOOPTS) $(pkgs)
+
+.PHONY: test-exp-short
+test-exp-short:
+	cd exp && $(GOTEST) -short $(GOOPTS) $(pkgs)

api/prometheus/v1/remote/remote_api.go → exp/api/remote/remote_api.go

@@ -22,23 +22,26 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"net/url"
+	"path"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/klauspost/compress/snappy"
 	"google.golang.org/protobuf/proto"
 
-	"github.com/prometheus/client_golang/api"
-	"github.com/prometheus/client_golang/internal/github.com/efficientgo/core/backoff"
+	"github.com/prometheus/client_golang/exp/internal/github.com/efficientgo/core/backoff"
 )
 
 // API is a client for Prometheus Remote Protocols.
 // NOTE(bwplotka): Only https://prometheus.io/docs/specs/remote_write_spec_2_0/ is currently implemented,
 // read protocols to be implemented if there will be a demand.
 type API struct {
-	client api.Client
-	opts   apiOpts
+	baseURL *url.URL
+	client  *http.Client
+
+	opts apiOpts
 
 	reqBuf, comprBuf []byte
 }
@@ -92,6 +95,14 @@ func WithAPINoRetryOnRateLimit() APIOption {
 	}
 }
 
+// WithAPIBackoff returns APIOption that allows overriding backoff configuration.
+func WithAPIBackoff(backoff backoff.Config) APIOption {
+	return func(o *apiOpts) error {
+		o.backoff = backoff
+		return nil
+	}
+}
+
 type nopSlogHandler struct{}
 
 func (n nopSlogHandler) Enabled(context.Context, slog.Level) bool  { return false }
@@ -103,7 +114,12 @@ func (n nopSlogHandler) WithGroup(string) slog.Handler             { return n }
 //
 // It is not safe to use the returned API from multiple goroutines, create a
 // separate *API for each goroutine.
-func NewAPI(c api.Client, opts ...APIOption) (*API, error) {
+func NewAPI(client *http.Client, baseURL string, opts ...APIOption) (*API, error) {
+	parsedURL, err := url.Parse(baseURL)
+	if err != nil {
+		return nil, fmt.Errorf("invalid base URL: %w", err)
+	}
+
 	o := *defaultAPIOpts
 	for _, opt := range opts {
 		if err := opt(&o); err != nil {
@@ -115,9 +131,16 @@ func NewAPI(c api.Client, opts ...APIOption) (*API, error) {
 		o.logger = slog.New(nopSlogHandler{})
 	}
 
+	if client == nil {
+		client = http.DefaultClient
+	}
+
+	parsedURL.Path = path.Join(parsedURL.Path, o.path)
+
 	return &API{
-		client: c,
-		opts:   o,
+		opts:    o,
+		client:  client,
+		baseURL: parsedURL,
 	}, nil
 }
 
@@ -157,6 +180,9 @@ type v2Request interface {
 //   - If neither is supported, it will marshaled using generic google.golang.org/protobuf methods and
 //     error out on unknown scheme.
 func (r *API) Write(ctx context.Context, msg any) (_ WriteResponseStats, err error) {
+	// Reset the buffer.
+	r.reqBuf = r.reqBuf[:0]
+
 	// Detect content-type.
 	cType := WriteProtoFullNameV1
 	if _, ok := msg.(v2Request); ok {
@@ -189,7 +215,6 @@ func (r *API) Write(ctx context.Context, msg any) (_ WriteResponseStats, err err
 		}
 	case proto.Message:
 		// Generic proto.
-		r.reqBuf = r.reqBuf[:0]
 		r.reqBuf, err = (proto.MarshalOptions{}).MarshalAppend(r.reqBuf, m)
 		if err != nil {
 			return WriteResponseStats{}, fmt.Errorf("encoding request %w", err)
@@ -266,8 +291,7 @@ func compressPayload(tmpbuf *[]byte, enc Compression, inp []byte) (compressed []
 }
 
 func (r *API) attemptWrite(ctx context.Context, compr Compression, proto WriteProtoFullName, payload []byte, attempt int) (WriteResponseStats, error) {
-	u := r.client.URL(r.opts.path, nil)
-	req, err := http.NewRequest(http.MethodPost, u.String(), bytes.NewReader(payload))
+	req, err := http.NewRequest(http.MethodPost, r.baseURL.String(), bytes.NewReader(payload))
 	if err != nil {
 		// Errors from NewRequest are from unparsable URLs, so are not
 		// recoverable.
@@ -287,11 +311,17 @@ func (r *API) attemptWrite(ctx context.Context, compr Compression, proto WritePr
 		req.Header.Set("Retry-Attempt", strconv.Itoa(attempt))
 	}
 
-	resp, body, err := r.client.Do(ctx, req)
+	resp, err := r.client.Do(req.WithContext(ctx))
 	if err != nil {
 		// Errors from Client.Do are likely network errors, so recoverable.
 		return WriteResponseStats{}, retryableError{err, 0}
 	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return WriteResponseStats{}, fmt.Errorf("reading response body: %w", err)
+	}
 
 	rs := WriteResponseStats{}
 	if proto == WriteProtoFullNameV2 {
@@ -334,19 +364,14 @@ type writeStorage interface {
 	Store(ctx context.Context, proto WriteProtoFullName, serializedRequest []byte) (_ WriteResponseStats, code int, _ error)
 }
 
-// remoteWriteDecompressor is an interface that allows decompressing the body of the request.
-type remoteWriteDecompressor interface {
-	Decompress(ctx context.Context, body io.ReadCloser) (decompressed []byte, _ error)
-}
-
 type handler struct {
 	store writeStorage
 	opts  handlerOpts
 }
 
 type handlerOpts struct {
-	logger       *slog.Logger
-	decompressor remoteWriteDecompressor
+	logger      *slog.Logger
+	middlewares []func(http.Handler) http.Handler
 }
 
 // HandlerOption represents an option for the handler.
@@ -360,22 +385,71 @@ func WithHandlerLogger(logger *slog.Logger) HandlerOption {
 	}
 }
 
-// WithHandlerDecompressor returns HandlerOption that allows providing remoteWriteDecompressor.
-// By default, SimpleSnappyDecompressor is used.
-func WithHandlerDecompressor(decompressor remoteWriteDecompressor) HandlerOption {
+// WithHandlerMiddleware returns HandlerOption that allows providing middlewares.
+// Multiple middlewares can be provided and will be applied in the order they are passed.
+// When using this option, SnappyDecompressorMiddleware is not applied by default so
+// it (or any other decompression middleware) needs to be added explicitly.
+func WithHandlerMiddlewares(middlewares ...func(http.Handler) http.Handler) HandlerOption {
 	return func(o *handlerOpts) {
-		o.decompressor = decompressor
+		o.middlewares = middlewares
+	}
+}
+
+// SnappyDecompressorMiddleware returns a middleware that checks if the request body is snappy-encoded and decompresses it.
+// If the request body is not snappy-encoded, it returns an error.
+// Used by default in NewRemoteWriteHandler.
+func SnappyDecompressorMiddleware(logger *slog.Logger) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			enc := r.Header.Get("Content-Encoding")
+			if enc != "" && enc != string(SnappyBlockCompression) {
+				err := fmt.Errorf("%v encoding (compression) is not accepted by this server; only %v is acceptable", enc, SnappyBlockCompression)
+				logger.Error("Error decoding remote write request", "err", err)
+				http.Error(w, err.Error(), http.StatusUnsupportedMediaType)
+				return
+			}
+
+			// Read the request body.
+			bodyBytes, err := io.ReadAll(r.Body)
+			if err != nil {
+				logger.Error("Error reading request body", "err", err)
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return
+			}
+
+			decompressed, err := snappy.Decode(nil, bodyBytes)
+			if err != nil {
+				// TODO(bwplotka): Add more context to responded error?
+				logger.Error("Error snappy decoding remote write request", "err", err)
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return
+			}
+
+			// Replace the body with decompressed data
+			r.Body = io.NopCloser(bytes.NewReader(decompressed))
+			next.ServeHTTP(w, r)
+		})
 	}
 }
 
 // NewRemoteWriteHandler returns HTTP handler that receives Remote Write 2.0
 // protocol https://prometheus.io/docs/specs/remote_write_spec_2_0/.
 func NewRemoteWriteHandler(store writeStorage, opts ...HandlerOption) http.Handler {
-	o := handlerOpts{logger: slog.New(nopSlogHandler{}), decompressor: &SimpleSnappyDecompressor{}}
+	o := handlerOpts{
+		logger:      slog.New(nopSlogHandler{}),
+		middlewares: []func(http.Handler) http.Handler{SnappyDecompressorMiddleware(slog.New(nopSlogHandler{}))},
+	}
 	for _, opt := range opts {
 		opt(&o)
 	}
-	return &handler{opts: o, store: store}
+	h := &handler{opts: o, store: store}
+
+	// Apply all middlewares in order
+	var handler http.Handler = h
+	for i := len(o.middlewares) - 1; i >= 0; i-- {
+		handler = o.middlewares[i](handler)
+	}
+	return handler
 }
 
 // ParseProtoMsg parses the content-type header and returns the proto message type.
@@ -412,10 +486,13 @@ func ParseProtoMsg(contentType string) (WriteProtoFullName, error) {
 }
 
 func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
 	contentType := r.Header.Get("Content-Type")
 	if contentType == "" {
-		// Don't break yolo 1.0 clients if not needed.
-		// We could give http.StatusUnsupportedMediaType, but let's assume 1.0 message by default.
 		contentType = appProtoContentType
 	}
 
@@ -426,26 +503,15 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	enc := r.Header.Get("Content-Encoding")
-	if enc == "" {
-		// Don't break yolo 1.0 clients if not needed. This is similar to what we did
-		// before 2.0: https://github.com/prometheus/prometheus/blob/d78253319daa62c8f28ed47e40bafcad2dd8b586/storage/remote/write_handler.go#L62
-		// We could give http.StatusUnsupportedMediaType, but let's assume snappy by default.
-	} else if enc != string(SnappyBlockCompression) {
-		err := fmt.Errorf("%v encoding (compression) is not accepted by this server; only %v is acceptable", enc, SnappyBlockCompression)
-		h.opts.logger.Error("Error decoding remote write request", "err", err)
-		http.Error(w, err.Error(), http.StatusUnsupportedMediaType)
-	}
-
-	// Decompress the request body.
-	decompressed, err := h.opts.decompressor.Decompress(r.Context(), r.Body)
+	// Read the already decompressed body
+	body, err := io.ReadAll(r.Body)
 	if err != nil {
-		h.opts.logger.Error("Error decompressing remote write request", "err", err.Error())
+		h.opts.logger.Error("Error reading request body", "err", err.Error())
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
 
-	stats, code, storeErr := h.store.Store(r.Context(), msgType, decompressed)
+	stats, code, storeErr := h.store.Store(r.Context(), msgType, body)
 
 	// Set required X-Prometheus-Remote-Write-Written-* response headers, in all cases.
 	stats.SetHeaders(w)
@@ -454,32 +520,11 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		if code == 0 {
 			code = http.StatusInternalServerError
 		}
-		if code/5 == 100 { // 5xx
+		if code/100 == 5 { // 5xx
 			h.opts.logger.Error("Error while storing the remote write request", "err", storeErr.Error())
 		}
 		http.Error(w, storeErr.Error(), code)
 		return
 	}
 	w.WriteHeader(http.StatusNoContent)
 }
-
-// SimpleSnappyDecompressor is a simple implementation of the remoteWriteDecompressor interface.
-type SimpleSnappyDecompressor struct{}
-
-func (s *SimpleSnappyDecompressor) Decompress(ctx context.Context, body io.ReadCloser) (decompressed []byte, _ error) {
-	// Read the request body.
-	bodyBytes, err := io.ReadAll(body)
-	if err != nil {
-		return nil, fmt.Errorf("error reading request body: %w", err)
-	}
-
-	decompressed, err = snappy.Decode(nil, bodyBytes)
-	if err != nil {
-		// TODO(bwplotka): Add more context to responded error?
-		return nil, fmt.Errorf("error snappy decoding request body: %w", err)
-	}
-
-	return decompressed, nil
-}
-
-var _ remoteWriteDecompressor = &SimpleSnappyDecompressor{}