remote: simplified Store interface; renamed Handler to WriteHandler; encapsulated write response (#1855)

* Re store: We don't need to pass context - you have it in
http.Request.Context()
* Re rename: Technically we have API and Handler is only for remote
write, not read, thus making space for "read" one day.
* Re write response: I made read flow private, otherwise it's confusing
how to use this structure from outside.

Signed-off-by: bwplotka <[email protected]>

Author: bwplotka

exp/api/remote/remote_api.go

@@ -64,8 +64,7 @@ var defaultAPIOpts = &apiOpts{
 		Max:        10 * time.Second,
 		MaxRetries: 10,
 	},
-	client: http.DefaultClient,
-	// Hardcoded for now.
+	client:           http.DefaultClient,
 	retryOnRateLimit: true,
 	compression:      SnappyBlockCompression,
 	path:             "api/v1/write",
@@ -96,7 +95,7 @@ func WithAPIPath(path string) APIOption {
 	}
 }
 
-// WithAPIRetryOnRateLimit returns APIOption that disables retrying on rate limit status code.
+// WithAPINoRetryOnRateLimit returns APIOption that disables retrying on rate limit status code.
 func WithAPINoRetryOnRateLimit() APIOption {
 	return func(o *apiOpts) error {
 		o.retryOnRateLimit = false
@@ -373,45 +372,46 @@ type writeStorage interface {
 	// Other headers might be trimmed, depending on the configured middlewares
 	// e.g. a default SnappyMiddleware trims "Content-Encoding" and ensures that
 	// encoded body bytes are already decompressed.
-	Store(ctx context.Context, msgType WriteMessageType, req *http.Request) (_ *WriteResponse, _ error)
+	Store(req *http.Request, msgType WriteMessageType) (_ *WriteResponse, _ error)
 }
 
-type handler struct {
+type writeHandler struct {
 	store                writeStorage
 	acceptedMessageTypes MessageTypes
-	opts                 handlerOpts
+	opts                 writeHandlerOpts
 }
 
-type handlerOpts struct {
+type writeHandlerOpts struct {
 	logger      *slog.Logger
 	middlewares []func(http.Handler) http.Handler
 }
 
-// HandlerOption represents an option for the handler.
-type HandlerOption func(o *handlerOpts)
+// WriteHandlerOption represents an option for the write handler.
+type WriteHandlerOption func(o *writeHandlerOpts)
 
-// WithHandlerLogger returns HandlerOption that allows providing slog logger.
+// WithWriteHandlerLogger returns WriteHandlerOption that allows providing slog logger.
 // By default, nothing is logged.
-func WithHandlerLogger(logger *slog.Logger) HandlerOption {
-	return func(o *handlerOpts) {
+func WithWriteHandlerLogger(logger *slog.Logger) WriteHandlerOption {
+	return func(o *writeHandlerOpts) {
 		o.logger = logger
 	}
 }
 
-// WithHandlerMiddleware returns HandlerOption that allows providing middlewares.
+// WithWriteHandlerMiddlewares returns WriteHandlerOption 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) {
+// This option replaces the default middlewares (SnappyDecompressorMiddleware), so if
+// you want to have handler that works with the default Remote Write 2.0 protocol,
+// SnappyDecompressorMiddleware (or any other decompression middleware) needs to be added explicitly.
+func WithWriteHandlerMiddlewares(middlewares ...func(http.Handler) http.Handler) WriteHandlerOption {
+	return func(o *writeHandlerOpts) {
 		o.middlewares = middlewares
 	}
 }
 
-// SnappyDecompressorMiddleware returns a middleware that checks if the request body is snappy-encoded and decompresses it.
+// SnappyDecodeMiddleware 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 {
+// Used by default in NewHandler.
+func SnappyDecodeMiddleware(logger *slog.Logger) func(http.Handler) http.Handler {
 	bufPool := sync.Pool{
 		New: func() any {
 			return bytes.NewBuffer(nil)
@@ -455,18 +455,18 @@ func SnappyDecompressorMiddleware(logger *slog.Logger) func(http.Handler) http.H
 	}
 }
 
-// NewHandler returns HTTP handler that receives Remote Write 2.0
+// NewWriteHandler returns HTTP handler that receives Remote Write 2.0
 // protocol https://prometheus.io/docs/specs/remote_write_spec_2_0/.
-func NewHandler(store writeStorage, acceptedMessageTypes MessageTypes, opts ...HandlerOption) http.Handler {
-	o := handlerOpts{
+func NewWriteHandler(store writeStorage, acceptedMessageTypes MessageTypes, opts ...WriteHandlerOption) http.Handler {
+	o := writeHandlerOpts{
 		logger:      slog.New(nopSlogHandler{}),
-		middlewares: []func(http.Handler) http.Handler{SnappyDecompressorMiddleware(slog.New(nopSlogHandler{}))},
+		middlewares: []func(http.Handler) http.Handler{SnappyDecodeMiddleware(slog.New(nopSlogHandler{}))},
 	}
 	for _, opt := range opts {
 		opt(&o)
 	}
 
-	h := &handler{
+	h := &writeHandler{
 		opts:                 o,
 		store:                store,
 		acceptedMessageTypes: acceptedMessageTypes,
@@ -513,7 +513,7 @@ func ParseProtoMsg(contentType string) (WriteMessageType, error) {
 	return WriteV1MessageType, nil
 }
 
-func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+func (h *writeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPost {
 		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
 		return
@@ -538,20 +538,25 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	writeResponse, storeErr := h.store.Store(r.Context(), msgType, r)
+	writeResponse, storeErr := h.store.Store(r, msgType)
+	if writeResponse == nil {
+		// User could forget to return write response; in this case we assume 0 samples
+		// were written.
+		writeResponse = NewWriteResponse()
+	}
 
-	// Set required X-Prometheus-Remote-Write-Written-* response headers, in all cases, alongwith any user-defined headers.
-	writeResponse.SetHeaders(w)
+	// Set required X-Prometheus-Remote-Write-Written-* response headers, in all cases, along with any user-defined headers.
+	writeResponse.writeHeaders(w)
 
 	if storeErr != nil {
-		if writeResponse.StatusCode() == 0 {
+		if writeResponse.statusCode == 0 {
 			writeResponse.SetStatusCode(http.StatusInternalServerError)
 		}
-		if writeResponse.StatusCode()/100 == 5 { // 5xx
-			h.opts.logger.Error("Error while storing the remote write request", "err", storeErr.Error())
+		if writeResponse.statusCode/100 == 5 { // 5xx
+			h.opts.logger.Error("Error while storing the remote write request", "err", storeErr)
 		}
-		http.Error(w, storeErr.Error(), writeResponse.StatusCode())
+		http.Error(w, storeErr.Error(), writeResponse.statusCode)
 		return
 	}
-	w.WriteHeader(writeResponse.StatusCode())
+	w.WriteHeader(writeResponse.statusCode)
 }

exp/api/remote/remote_api_test.go

@@ -71,8 +71,8 @@ type mockStorage struct {
 	mockErr  error
 }
 
-func (m *mockStorage) Store(_ context.Context, msgType WriteMessageType, req *http.Request) (*WriteResponse, error) {
-	w := &WriteResponse{}
+func (m *mockStorage) Store(req *http.Request, msgType WriteMessageType) (*WriteResponse, error) {
+	w := NewWriteResponse()
 	if m.mockErr != nil {
 		if m.mockCode != nil {
 			w.SetStatusCode(*m.mockCode)
@@ -147,7 +147,7 @@ func TestRemoteAPI_Write_WithHandler(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		tLogger := slog.Default()
 		mStore := &mockStorage{}
-		srv := httptest.NewServer(NewHandler(mStore, MessageTypes{WriteV2MessageType}, WithHandlerLogger(tLogger)))
+		srv := httptest.NewServer(NewWriteHandler(mStore, MessageTypes{WriteV2MessageType}, WithWriteHandlerLogger(tLogger)))
 		t.Cleanup(srv.Close)
 
 		client, err := NewAPI(srv.URL,
@@ -182,7 +182,7 @@ func TestRemoteAPI_Write_WithHandler(t *testing.T) {
 			mockErr:  errors.New("storage error"),
 			mockCode: &mockCode,
 		}
-		srv := httptest.NewServer(NewHandler(mStore, MessageTypes{WriteV2MessageType}, WithHandlerLogger(tLogger)))
+		srv := httptest.NewServer(NewWriteHandler(mStore, MessageTypes{WriteV2MessageType}, WithWriteHandlerLogger(tLogger)))
 		t.Cleanup(srv.Close)
 
 		client, err := NewAPI(srv.URL,

exp/api/remote/remote_headers.go

@@ -125,29 +125,19 @@ func (w *WriteResponse) SetStatusCode(code int) {
 	w.statusCode = code
 }
 
-// StatusCode returns the current HTTP status code.
-func (w *WriteResponse) StatusCode() int {
-	return w.statusCode
-}
-
 // SetExtraHeader adds additional headers to be set in the response (apart from stats headers)
 func (w *WriteResponse) SetExtraHeader(key, value string) {
 	w.extraHeaders.Set(key, value)
 }
 
-// ExtraHeaders returns all additional headers to be set in the response (apart from stats headers).
-func (w *WriteResponse) ExtraHeaders() http.Header {
-	return w.extraHeaders
-}
-
-// SetHeaders sets response headers in a given response writer.
+// writeHeaders sets response headers in a given response writer.
 // Make sure to use it before http.ResponseWriter.WriteHeader and .Write.
-func (r *WriteResponse) SetHeaders(w http.ResponseWriter) {
-	h := w.Header()
-	h.Set(writtenSamplesHeader, strconv.Itoa(r.Samples))
-	h.Set(writtenHistogramsHeader, strconv.Itoa(r.Histograms))
-	h.Set(writtenExemplarsHeader, strconv.Itoa(r.Exemplars))
-	for k, v := range r.ExtraHeaders() {
+func (w *WriteResponse) writeHeaders(rw http.ResponseWriter) {
+	h := rw.Header()
+	h.Set(writtenSamplesHeader, strconv.Itoa(w.Samples))
+	h.Set(writtenHistogramsHeader, strconv.Itoa(w.Histograms))
+	h.Set(writtenExemplarsHeader, strconv.Itoa(w.Exemplars))
+	for k, v := range w.extraHeaders {
 		for _, vv := range v {
 			h.Add(k, vv)
 		}

exp/api/remote/remote_headers_test.go

@@ -25,16 +25,16 @@ import (
 func TestWriteResponse(t *testing.T) {
 	t.Run("new response has empty headers", func(t *testing.T) {
 		resp := NewWriteResponse()
-		if len(resp.ExtraHeaders()) != 0 {
-			t.Errorf("expected empty headers, got %v", resp.ExtraHeaders())
+		if len(resp.extraHeaders) != 0 {
+			t.Errorf("expected empty headers, got %v", resp.extraHeaders)
 		}
 	})
 
-	t.Run("setters and getters", func(t *testing.T) {
+	t.Run("setters", func(t *testing.T) {
 		resp := NewWriteResponse()
 
 		resp.SetStatusCode(http.StatusOK)
-		if got := resp.StatusCode(); got != http.StatusOK {
+		if got := resp.statusCode; got != http.StatusOK {
 			t.Errorf("expected status code %d, got %d", http.StatusOK, got)
 		}
 
@@ -66,12 +66,12 @@ func TestWriteResponse(t *testing.T) {
 		}
 
 		resp.SetExtraHeader("Test-Header", "test-value")
-		if got := resp.ExtraHeaders().Get("Test-Header"); got != "test-value" {
+		if got := resp.extraHeaders.Get("Test-Header"); got != "test-value" {
 			t.Errorf("expected header value %q, got %q", "test-value", got)
 		}
 	})
 
-	t.Run("set headers on response writer", func(t *testing.T) {
+	t.Run("writeHeaders", func(t *testing.T) {
 		resp := NewWriteResponse()
 		resp.Add(WriteResponseStats{
 			Samples:    10,
@@ -82,7 +82,7 @@ func TestWriteResponse(t *testing.T) {
 		resp.SetExtraHeader("Custom-Header", "custom-value")
 
 		w := httptest.NewRecorder()
-		resp.SetHeaders(w)
+		resp.writeHeaders(w)
 
 		expectedHeaders := map[string]string{
 			"Custom-Header": "custom-value",