feat: Enable warnings in Loki query responses #12425

MasslessParticle · 2024-04-01T22:14:49Z

Under certain conditions, it might be useful for Loki to be able to return warnings to the frontend. The prometheus response types support this so it's just a matter of wiring it in.

MasslessParticle · 2024-04-01T22:17:22Z

pkg/ingester/ingester.go

@@ -874,6 +875,7 @@ func (i *Ingester) GetOrCreateInstance(instanceID string) (*instance, error) { /
 func (i *Ingester) Query(req *logproto.QueryRequest, queryServer logproto.Querier_QueryServer) error {
 	// initialize stats collection for ingester queries.
 	_, ctx := stats.NewContext(queryServer.Context())
+	_, ctx = metadata.NewContext(ctx)


This creates a new metadata context for the ingester. Any warnings will be attached in the returned batch stats.

MasslessParticle · 2024-04-01T22:18:25Z

pkg/ingester/instance.go

@@ -971,6 +973,7 @@ func sendBatches(ctx context.Context, i iter.EntryIterator, queryServer QuerierQ

 		stats.AddIngesterBatch(int64(batchSize))
 		batch.Stats = stats.Ingester()
+		batch.Warnings = metadata.Warnings()


Here's where we attach warnings to the ingester batch response. From here on out, it's a matter of encoding and decoding (and encoding and decoding) in the scheduler and frontend to propagate them to the user

MasslessParticle · 2024-04-01T22:21:04Z

pkg/logql/engine.go

@@ -266,6 +266,7 @@ func (q *query) Exec(ctx context.Context) (logqlmodel.Result, error) {
 		Data:       data,
 		Statistics: statResult,
 		Headers:    metadataCtx.Headers(),
+		Warnings:   metadataCtx.Warnings(),


This is the crux of the PR on the frontend. Everything comes through here so this is where we attach the warnings to be propagated to the scheduler and querier.

MasslessParticle · 2024-04-01T22:22:00Z

pkg/querier/queryrange/codec.go

@@ -1145,7 +1144,8 @@ func decodeResponseJSONFrom(buf []byte, req queryrangebase.Request, headers http
 						ResultType: loghttp.ResultTypeMatrix,
 						Result:     toProtoMatrix(resp.Data.Result.(loghttp.Matrix)),
 					},
-					Headers: convertPrometheusResponseHeadersToPointers(httpResponseHeadersToPromResponseHeaders(headers)),
+					Headers:  convertPrometheusResponseHeadersToPointers(httpResponseHeadersToPromResponseHeaders(headers)),


In this file, we have to attach warnings to each of the response types that might possibly have them

MasslessParticle · 2024-04-01T22:23:25Z

pkg/querier/queryrange/codec.go

 	for _, res := range responses {
 		lokiResult := res.(*LokiResponse)
 		mergedStats.MergeSplit(lokiResult.Statistics)
 		lokiResponses = append(lokiResponses, lokiResult)
+


Loki splits a query and then merges the responses into one. Here we make sure all the warnings from subqueries make it into the merged query. We'll do the same below for metrics queries

MasslessParticle · 2024-04-01T22:23:52Z

pkg/querier/queryrange/queryrangebase/query_range.go

@@ -164,8 +164,28 @@ func (p prometheusCodec) MergeResponse(responses ...Response) (Response, error)
 	// Merge the responses.
 	sort.Sort(byFirstTime(promResponses))

+	uniqueWarnings := map[string]struct{}{}


Merging metrics queries

MasslessParticle · 2024-04-01T22:24:13Z

pkg/util/marshal/query.go

 	s.WriteObjectStart()
 	s.WriteObjectField("status")
 	s.WriteString("success")

+	if len(warnings) > 0 {


Custom json marshalling

MasslessParticle · 2024-04-01T22:24:41Z

pkg/loghttp/query.go

 }

 func (q *QueryResponse) UnmarshalJSON(data []byte) error {
 	return jsonparser.ObjectEach(data, func(key, value []byte, dataType jsonparser.ValueType, offset int) error {
 		switch string(key) {
 		case "status":
 			q.Status = string(value)
+		case "warnings":


Custom json unmarshalling

na-- · 2024-04-02T14:41:24Z

pkg/logqlmodel/metadata/context.go

@@ -94,3 +119,19 @@ func ExtendHeaders(dst map[string][]string, src []*definitions.PrometheusRespons
 		dst[header.Name] = header.Values
 	}
 }
+
+func AddWarnings(ctx context.Context, warnings ...string) error {
+	context, ok := ctx.Value(metadataKey).(*Context)


Presumably, most of the time there will be no warnings, and looking up keys in contexts can be surprisingly expensive, not to mention the locking, so a short circuit here is probably a good idea 🤔

Suggested change

context, ok := ctx.Value(metadataKey).(*Context)

if len(warnings) == 0 {

return nil

}

context, ok := ctx.Value(metadataKey).(*Context)

na-- · 2024-04-02T14:53:56Z

pkg/querier/queryrange/codec.go

+		}
+	}
+
+	warnings := make([]string, 0, len(uniqueWarnings))


Nit, but since we already import golang.org/x/exp, we might as well use maps.Keys from it

na-- · 2024-04-02T14:55:15Z

pkg/logqlmodel/metadata/context.go

+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+
+	warnings := make([]string, 0, len(c.warnings))


Same as https://github.com/grafana/loki/pull/12425/files#r1548053553, we can use maps.Keys() here

na-- · 2024-04-02T14:59:07Z

pkg/logql/accumulator.go

@@ -353,6 +355,14 @@ func (acc *AccumulatedStreams) Result() []logqlmodel.Result {
 		)
 	}

+	warnings := make([]string, 0, len(acc.warnings))


Same, maps.Keys() can simplify this

na-- · 2024-04-02T15:00:44Z

pkg/querier/queryrange/queryrangebase/query_range.go

+		}
+	}
+
+	warnings := make([]string, 0, len(uniqueWarnings))


Same, maps.Keys() 😅

na-- · 2024-04-02T15:02:29Z

pkg/querier/queryrange/querysharding.go

@@ -252,6 +252,7 @@ func (ast *astMapperware) Do(ctx context.Context, r queryrangebase.Request) (que
 		return nil, err
 	}

+	// TODO(tpatterson): add warnings to response


This TODO seems obsolete, you seem to have added the warnings to all of the response types below ?

na-- · 2024-04-02T15:07:30Z

pkg/loghttp/query.go

+		case "warnings":
+			var warnings []string
+			if _, err := jsonparser.ArrayEach(value, func(value []byte, dataType jsonparser.ValueType, offset int, err error) {
+				warnings = append(warnings, string(value))


Hmm considering JSON can contain all sorts of junk in an array, this probably needs a bit more validation that we are actually getting strings here, presumably through the dataType 🤔

I assume that value doesn't contain the string quote characters, but I wonder if it contains escape characters, e.g. if the warning itself contains a quote symbol, I assume it will be escaped in value here?

na--

I don't fully understand how every part of this is threaded and fits together, but overall the code LGTM, with the minor nitpicks I added as inline comment.

I'll 👍 to unblock, though the only inline comment that is somewhat blocking is this one, though this one is probably also quite nice to have, considering how AddWarnings() seems to be very much on the hot path.

feat: Enable warnings in Loki query responses

f2b3963

MasslessParticle requested a review from a team as a code owner April 1, 2024 22:14

pull-request-size bot added the size/L label Apr 1, 2024

MasslessParticle commented Apr 1, 2024

View reviewed changes

MasslessParticle added 2 commits April 1, 2024 16:19

remove debug lines

822ea1c

lint

e44351e

MasslessParticle commented Apr 1, 2024

View reviewed changes

na-- reviewed Apr 2, 2024

View reviewed changes

na-- approved these changes Apr 2, 2024

View reviewed changes

MasslessParticle added 5 commits April 2, 2024 15:23

Merge branch 'main' into tpatterson/response-warnings

5111860

Merge branch 'main' into tpatterson/response-warnings

d689e1f

Merge branch 'main' into tpatterson/response-warnings

cab8da2

review feedback

d302464

lint

aca94ff

MasslessParticle merged commit cf04ec1 into main Apr 3, 2024
10 checks passed

MasslessParticle deleted the tpatterson/response-warnings branch April 3, 2024 16:42

rhnasc pushed a commit to inloco/loki that referenced this pull request Apr 12, 2024

feat: Enable warnings in Loki query responses (grafana#12425)

150f8a4

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Enable warnings in Loki query responses #12425

feat: Enable warnings in Loki query responses #12425

MasslessParticle commented Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

MasslessParticle Apr 1, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- Apr 2, 2024

na-- left a comment

feat: Enable warnings in Loki query responses #12425

feat: Enable warnings in Loki query responses #12425

Conversation

MasslessParticle commented Apr 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

na-- left a comment

Choose a reason for hiding this comment