Make Write message type more flexble, address some feedback #1710
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -51,6 +51,7 @@ type apiOpts struct { | |
| logger *slog.Logger | ||
| backoff backoff.Config | ||
| compression Compression | ||
| path string | ||
| retryOnRateLimit bool | ||
| } | ||
|
|
||
|
|
@@ -63,6 +64,7 @@ var defaultAPIOpts = &apiOpts{ | |
| // Hardcoded for now. | ||
| retryOnRateLimit: true, | ||
| compression: SnappyBlockCompression, | ||
| path: "api/v1/write", | ||
| } | ||
|
|
||
| // WithAPILogger returns APIOption that allows providing slog logger. | ||
|
|
@@ -74,6 +76,22 @@ func WithAPILogger(logger *slog.Logger) APIOption { | |
| } | ||
| } | ||
|
|
||
| // WithAPIPath returns APIOption that allows providing path to send remote write requests to. | ||
| func WithAPIPath(path string) APIOption { | ||
| return func(o *apiOpts) error { | ||
| o.path = path | ||
| return nil | ||
| } | ||
| } | ||
|
|
||
| // WithAPIRetryOnRateLimit returns APIOption that disables retrying on rate limit status code. | ||
| func WithAPINoRetryOnRateLimit() APIOption { | ||
| return func(o *apiOpts) error { | ||
| o.retryOnRateLimit = false | ||
| return nil | ||
| } | ||
| } | ||
|
|
||
| type nopSlogHandler struct{} | ||
|
|
||
| func (n nopSlogHandler) Enabled(context.Context, slog.Level) bool { return false } | ||
|
|
@@ -117,33 +135,67 @@ type vtProtoEnabled interface { | |
| MarshalToSizedBufferVT(dAtA []byte) (int, error) | ||
| } | ||
|
|
||
| type gogoProtoEnabled interface { | ||
| Size() (n int) | ||
| MarshalToSizedBuffer(dAtA []byte) (n int, err error) | ||
| } | ||
|
|
||
| // Sort of a hack to identify v2 requests. | ||
| // Under any marshaling scheme, v2 requests have a `Symbols` field of type []string. | ||
| // So would always have a `GetSymbols()` method which doesn't rely on any other types. | ||
| type v2Request interface { | ||
| GetSymbols() []string | ||
| } | ||
|
|
||
| // Write writes given, non-empty, protobuf message to a remote storage. | ||
| // | ||
| // Depending on serialization methods, | ||
| // - https://github.com/planetscale/vtprotobuf methods will be used if your msg | ||
| // supports those (e.g. SizeVT() and MarshalToSizedBufferVT(...)), for efficiency | ||
| // - Otherwise https://github.com/gogo/protobuf methods (e.g. Size() and MarshalToSizedBuffer(...)) | ||
| // will be used | ||
| // - 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) { | ||
| // Detect content-type. | ||
| cType := WriteProtoFullNameV1 | ||
| if _, ok := msg.(v2Request); ok { | ||
| cType = WriteProtoFullNameV2 | ||
| } | ||
|
|
||
| if err := cType.Validate(); err != nil { | ||
| return WriteResponseStats{}, err | ||
| } | ||
|
|
||
| // Encode the payload. | ||
| switch m := msg.(type) { | ||
| case vtProtoEnabled: | ||
| // Use optimized vtprotobuf if supported. | ||
| size := m.SizeVT() | ||
| if len(r.reqBuf) < size { | ||
| r.reqBuf = make([]byte, size) | ||
| } | ||
| if _, err := m.MarshalToSizedBufferVT(r.reqBuf[:size]); err != nil { | ||
| return WriteResponseStats{}, fmt.Errorf("encoding request %w", err) | ||
| } | ||
| case gogoProtoEnabled: | ||
| // Gogo proto if supported. | ||
| size := m.Size() | ||
| if len(r.reqBuf) < size { | ||
| r.reqBuf = make([]byte, size) | ||
| } | ||
| if _, err := m.MarshalToSizedBuffer(r.reqBuf[:size]); err != nil { | ||
| return WriteResponseStats{}, fmt.Errorf("encoding request %w", err) | ||
| } | ||
| case proto.Message: | ||
saswatamcode marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| // 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) | ||
| } | ||
| default: | ||
| return WriteResponseStats{}, fmt.Errorf("unknown message type %T", m) | ||
| } | ||
|
|
||
| payload, err := compressPayload(&r.comprBuf, r.opts.compression, r.reqBuf) | ||
|
|
@@ -214,7 +266,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)) | ||
| if err != nil { | ||
| // Errors from NewRequest are from unparsable URLs, so are not | ||
|
|
@@ -241,9 +293,12 @@ func (r *API) attemptWrite(ctx context.Context, compr Compression, proto WritePr | |
| return WriteResponseStats{}, retryableError{err, 0} | ||
| } | ||
|
|
||
| rs := WriteResponseStats{} | ||
| if proto == WriteProtoFullNameV2 { | ||
| rs, err = parseWriteResponseStats(resp) | ||
| if err != nil { | ||
| r.opts.logger.Warn("parsing rw write statistics failed; partial or no stats", "err", err) | ||
| } | ||
| } | ||
|
|
||
| if resp.StatusCode/100 == 2 { | ||
|
|
@@ -279,18 +334,59 @@ 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 { | ||
|
There was a problem hiding this comment. can you remind me why we need this? Another option is to use server middleware interface for this. There are common middlewares for this work ppl use e.g. Otel There was a problem hiding this comment. Some downstream users have their own request body size limiter+instrumentaion impls, that they would probably want to use. I know Thanos and Cortex have them #1658 (comment) I don't think compression-wise things would be too different (maybe using s2). There was a problem hiding this comment. Cool, I will merge for now, but we need to switch to middleware injection technique |
||
| 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 | ||
| } | ||
|
|
||
| // HandlerOption represents an option for the handler. | ||
| type HandlerOption func(o *handlerOpts) | ||
|
|
||
| // WithHandlerLogger returns HandlerOption that allows providing slog logger. | ||
| // By default, nothing is logged. | ||
| func WithHandlerLogger(logger *slog.Logger) HandlerOption { | ||
| return func(o *handlerOpts) { | ||
| o.logger = logger | ||
| } | ||
| } | ||
|
|
||
| // WithHandlerDecompressor returns HandlerOption that allows providing remoteWriteDecompressor. | ||
| // By default, SimpleSnappyDecompressor is used. | ||
| func WithHandlerDecompressor(decompressor remoteWriteDecompressor) HandlerOption { | ||
| return func(o *handlerOpts) { | ||
| o.decompressor = decompressor | ||
| } | ||
| } | ||
|
|
||
| // 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{}} | ||
| for _, opt := range opts { | ||
| opt(&o) | ||
| } | ||
| return &handler{opts: o, store: store} | ||
| } | ||
|
|
||
| // ParseProtoMsg parses the content-type header and returns the proto message type. | ||
| // | ||
| // The expected content-type will be of the form, | ||
| // - `application/x-protobuf;proto=io.prometheus.write.v2.Request` which will be treated as RW2.0 request, | ||
| // - `application/x-protobuf;proto=prometheus.WriteRequest` which will be treated as RW1.0 request, | ||
| // - `application/x-protobuf` which will be treated as RW1.0 request. | ||
| // | ||
| // If the content-type is not of the above forms, it will return an error. | ||
| func ParseProtoMsg(contentType string) (WriteProtoFullName, error) { | ||
| contentType = strings.TrimSpace(contentType) | ||
|
|
||
| parts := strings.Split(contentType, ";") | ||
|
|
@@ -323,9 +419,9 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { | |
| contentType = appProtoContentType | ||
| } | ||
|
|
||
| msgType, err := ParseProtoMsg(contentType) | ||
| if err != nil { | ||
| h.opts.logger.Error("Error decoding remote write request", "err", err) | ||
| http.Error(w, err.Error(), http.StatusUnsupportedMediaType) | ||
| return | ||
| } | ||
|
|
@@ -337,22 +433,14 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { | |
| // 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) | ||
| if err != nil { | ||
| h.opts.logger.Error("Error decompressing remote write request", "err", err.Error()) | ||
| http.Error(w, err.Error(), http.StatusBadRequest) | ||
| return | ||
| } | ||
|
|
@@ -367,10 +455,31 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { | |
| code = http.StatusInternalServerError | ||
| } | ||
| if code/5 == 100 { // 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{} | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This seems a bit hacky, but need it for the contentType detection. I wonder if there is a better way here
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Interesting. I think this is acceptable for now.