Skip to content

Commit b272573

Browse files
bentonamcstyan
andauthored
feat: Add Rule Name and Type to Query Tags (#15055)
Signed-off-by: Callum Styan <callumstyan@gmail.com> Co-authored-by: Callum Styan <callumstyan@gmail.com>
1 parent d2b6fd8 commit b272573

5 files changed

+95
-12
lines changed

‎pkg/ruler/compat.go

+28
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,12 @@ func queryFunc(evaluator Evaluator, checker readyChecker, userID string, logger
7272
return nil, errNotReady
7373
}
7474

75+
// Extract rule details
76+
ruleName := detail.Name
77+
ruleType := detail.Kind
78+
79+
// Add rule details to context
80+
ctx = AddRuleDetailsToContext(ctx, ruleName, ruleType)
7581
res, err := evaluator.Eval(ctx, qs, t)
7682

7783
if err != nil {
@@ -357,3 +363,25 @@ type noopRuleDependencyController struct{}
357363
func (*noopRuleDependencyController) AnalyseRules([]rules.Rule) {
358364
// Do nothing
359365
}
366+
367+
// Define context keys to avoid collisions
368+
type contextKey string
369+
370+
const (
371+
ruleNameKey contextKey = "rule_name"
372+
ruleTypeKey contextKey = "rule_type"
373+
)
374+
375+
// AddRuleDetailsToContext adds rule details to the context
376+
func AddRuleDetailsToContext(ctx context.Context, ruleName string, ruleType string) context.Context {
377+
ctx = context.WithValue(ctx, ruleNameKey, ruleName)
378+
ctx = context.WithValue(ctx, ruleTypeKey, ruleType)
379+
return ctx
380+
}
381+
382+
// GetRuleDetailsFromContext retrieves rule details from the context
383+
func GetRuleDetailsFromContext(ctx context.Context) (string, string) {
384+
ruleName, _ := ctx.Value(ruleNameKey).(string)
385+
ruleType, _ := ctx.Value(ruleTypeKey).(string)
386+
return ruleName, ruleType
387+
}

‎pkg/ruler/compat_test.go

+16
Original file line numberDiff line numberDiff line change
@@ -130,3 +130,19 @@ type fakeChecker struct{}
130130
func (f fakeChecker) isReady(_ string) bool {
131131
return true
132132
}
133+
134+
func TestAddAndGetRuleDetailsFromContext(t *testing.T) {
135+
ctx := context.Background()
136+
ruleName := "test_rule"
137+
ruleType := "test_type"
138+
139+
// Add rule details to context
140+
ctx = AddRuleDetailsToContext(ctx, ruleName, ruleType)
141+
142+
// Retrieve rule details from context
143+
retrievedRuleName, retrievedRuleType := GetRuleDetailsFromContext(ctx)
144+
145+
// Assert that the retrieved values match the expected values
146+
assert.Equal(t, ruleName, retrievedRuleName, "Expected rule name to match")
147+
assert.Equal(t, ruleType, retrievedRuleType, "Expected rule type to match")
148+
}

‎pkg/ruler/evaluator_local.go

+13-1
Original file line numberDiff line numberDiff line change
@@ -3,28 +3,36 @@ package ruler
33
import (
44
"context"
55
"fmt"
6+
"os"
67
"time"
78

89
"github.com/go-kit/log"
10+
"github.com/go-kit/log/level"
911

1012
"github.com/grafana/loki/v3/pkg/logproto"
1113
"github.com/grafana/loki/v3/pkg/logql"
1214
"github.com/grafana/loki/v3/pkg/logqlmodel"
15+
"github.com/grafana/loki/v3/pkg/util"
1316
)
1417

1518
const EvalModeLocal = "local"
1619

1720
type LocalEvaluator struct {
1821
engine *logql.Engine
1922
logger log.Logger
23+
24+
// we don't want/need to log all the additional context, such as
25+
// caller=spanlogger.go:116 component=ruler evaluation_mode=remote method=ruler.remoteEvaluation.Query
26+
// in insights logs, so create a new logger
27+
insightsLogger log.Logger
2028
}
2129

2230
func NewLocalEvaluator(engine *logql.Engine, logger log.Logger) (*LocalEvaluator, error) {
2331
if engine == nil {
2432
return nil, fmt.Errorf("given engine is nil")
2533
}
2634

27-
return &LocalEvaluator{engine: engine, logger: logger}, nil
35+
return &LocalEvaluator{engine: engine, logger: logger, insightsLogger: log.NewLogfmtLogger(os.Stderr)}, nil
2836
}
2937

3038
func (l *LocalEvaluator) Eval(ctx context.Context, qs string, now time.Time) (*logqlmodel.Result, error) {
@@ -49,5 +57,9 @@ func (l *LocalEvaluator) Eval(ctx context.Context, qs string, now time.Time) (*l
4957
return nil, err
5058
}
5159

60+
// Retrieve rule details from context
61+
ruleName, ruleType := GetRuleDetailsFromContext(ctx)
62+
63+
level.Info(l.insightsLogger).Log("msg", "request timings", "insight", "true", "source", "loki_ruler", "rule_name", ruleName, "rule_type", ruleType, "total", res.Statistics.Summary.ExecTime, "total_bytes", res.Statistics.Summary.TotalBytesProcessed, "query_hash", util.HashedQuery(qs))
5264
return &res, nil
5365
}

‎pkg/ruler/evaluator_remote.go

+28-11
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"net/http"
1616
"net/textproto"
1717
"net/url"
18+
"os"
1819
"strconv"
1920
"time"
2021

@@ -72,15 +73,21 @@ type RemoteEvaluator struct {
7273
overrides RulesLimits
7374
logger log.Logger
7475

76+
// we don't want/need to log all the additional context, such as
77+
// caller=spanlogger.go:116 component=ruler evaluation_mode=remote method=ruler.remoteEvaluation.Query
78+
// in insights logs, so create a new logger
79+
insightsLogger log.Logger
80+
7581
metrics *metrics
7682
}
7783

7884
func NewRemoteEvaluator(client httpgrpc.HTTPClient, overrides RulesLimits, logger log.Logger, registerer prometheus.Registerer) (*RemoteEvaluator, error) {
7985
return &RemoteEvaluator{
80-
client: client,
81-
overrides: overrides,
82-
logger: logger,
83-
metrics: newMetrics(registerer),
86+
client: client,
87+
overrides: overrides,
88+
logger: logger,
89+
insightsLogger: log.NewLogfmtLogger(os.Stderr),
90+
metrics: newMetrics(registerer),
8491
}, nil
8592
}
8693

@@ -220,6 +227,11 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
220227
body := []byte(args.Encode())
221228
hash := util.HashedQuery(query)
222229

230+
// Retrieve rule details from context
231+
ruleName, ruleType := GetRuleDetailsFromContext(ctx)
232+
233+
// Construct the X-Query-Tags header value
234+
queryTags := fmt.Sprintf("source=ruler,rule_name=%s,rule_type=%s", ruleName, ruleType)
223235
req := httpgrpc.HTTPRequest{
224236
Method: http.MethodPost,
225237
Url: queryEndpointPath,
@@ -228,7 +240,7 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
228240
{Key: textproto.CanonicalMIMEHeaderKey("User-Agent"), Values: []string{userAgent}},
229241
{Key: textproto.CanonicalMIMEHeaderKey("Content-Type"), Values: []string{mimeTypeFormPost}},
230242
{Key: textproto.CanonicalMIMEHeaderKey("Content-Length"), Values: []string{strconv.Itoa(len(body))}},
231-
{Key: textproto.CanonicalMIMEHeaderKey(string(httpreq.QueryTagsHTTPHeader)), Values: []string{"source=ruler"}},
243+
{Key: textproto.CanonicalMIMEHeaderKey(string(httpreq.QueryTagsHTTPHeader)), Values: []string{queryTags}},
232244
{Key: textproto.CanonicalMIMEHeaderKey(user.OrgIDHeaderName), Values: []string{orgID}},
233245
},
234246
}
@@ -242,12 +254,12 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
242254
instrument.ObserveWithExemplar(ctx, r.metrics.responseSizeBytes.WithLabelValues(orgID), float64(len(resp.Body)))
243255
}
244256

245-
log := log.With(logger, "query_hash", hash, "query", query, "instant", ts, "response_time", time.Since(start).String())
257+
logger = log.With(logger, "query_hash", hash, "query", query, "instant", ts, "response_time", time.Since(start).String())
246258

247259
if err != nil {
248260
r.metrics.failedEvals.WithLabelValues("error", orgID).Inc()
249261

250-
level.Warn(log).Log("msg", "failed to evaluate rule", "err", err)
262+
level.Warn(logger).Log("msg", "failed to evaluate rule", "err", err)
251263
return nil, fmt.Errorf("rule evaluation failed: %w", err)
252264
}
253265

@@ -261,22 +273,27 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
261273
r.metrics.failedEvals.WithLabelValues("upstream_error", orgID).Inc()
262274

263275
respBod, _ := io.ReadAll(limitedBody)
264-
level.Warn(log).Log("msg", "rule evaluation failed with non-2xx response", "response_code", resp.Code, "response_body", respBod)
276+
level.Warn(logger).Log("msg", "rule evaluation failed with non-2xx response", "response_code", resp.Code, "response_body", respBod)
265277
return nil, fmt.Errorf("unsuccessful/unexpected response - status code %d", resp.Code)
266278
}
267279

268280
maxSize := r.overrides.RulerRemoteEvaluationMaxResponseSize(orgID)
269281
if maxSize > 0 && int64(len(fullBody)) >= maxSize {
270282
r.metrics.failedEvals.WithLabelValues("max_size", orgID).Inc()
271283

272-
level.Error(log).Log("msg", "rule evaluation exceeded max size", "max_size", maxSize, "response_size", len(fullBody))
284+
level.Error(logger).Log("msg", "rule evaluation exceeded max size", "max_size", maxSize, "response_size", len(fullBody))
273285
return nil, fmt.Errorf("%d bytes exceeds response size limit of %d (defined by ruler_remote_evaluation_max_response_size)", len(resp.Body), maxSize)
274286
}
275287

276-
level.Debug(log).Log("msg", "rule evaluation succeeded")
288+
level.Debug(logger).Log("msg", "rule evaluation succeeded")
277289
r.metrics.successfulEvals.WithLabelValues(orgID).Inc()
278290

279-
return r.decodeResponse(ctx, resp, orgID)
291+
dr, err := r.decodeResponse(ctx, resp, orgID)
292+
if err != nil {
293+
return nil, err
294+
}
295+
level.Info(r.insightsLogger).Log("msg", "request timings", "insight", "true", "source", "loki_ruler", "rule_name", ruleName, "rule_type", ruleType, "total", dr.Statistics.Summary.ExecTime, "total_bytes", dr.Statistics.Summary.TotalBytesProcessed, "query_hash", util.HashedQuery(query))
296+
return dr, err
280297
}
281298

282299
func (r *RemoteEvaluator) decodeResponse(ctx context.Context, resp *httpgrpc.HTTPResponse, orgID string) (*logqlmodel.Result, error) {

‎pkg/ruler/evaluator_remote_test.go

+10
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ func TestRemoteEvalQueryTimeout(t *testing.T) {
6161

6262
ctx := context.Background()
6363
ctx = user.InjectOrgID(ctx, "test")
64+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
6465

6566
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
6667
require.Error(t, err)
@@ -98,6 +99,7 @@ func TestRemoteEvalMaxResponseSize(t *testing.T) {
9899

99100
ctx := context.Background()
100101
ctx = user.InjectOrgID(ctx, "test")
102+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
101103

102104
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
103105
require.Error(t, err)
@@ -146,6 +148,7 @@ func TestRemoteEvalScalar(t *testing.T) {
146148

147149
ctx := context.Background()
148150
ctx = user.InjectOrgID(ctx, "test")
151+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
149152

150153
res, err := ev.Eval(ctx, "19", now)
151154
require.NoError(t, err)
@@ -189,6 +192,7 @@ func TestRemoteEvalEmptyScalarResponse(t *testing.T) {
189192

190193
ctx := context.Background()
191194
ctx = user.InjectOrgID(ctx, "test")
195+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
192196

193197
res, err := ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
194198
require.NoError(t, err)
@@ -247,6 +251,7 @@ func TestRemoteEvalVectorResponse(t *testing.T) {
247251

248252
ctx := context.Background()
249253
ctx = user.InjectOrgID(ctx, "test")
254+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
250255

251256
res, err := ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", now)
252257
require.NoError(t, err)
@@ -294,6 +299,7 @@ func TestRemoteEvalEmptyVectorResponse(t *testing.T) {
294299

295300
ctx := context.Background()
296301
ctx = user.InjectOrgID(ctx, "test")
302+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
297303

298304
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
299305
require.NoError(t, err)
@@ -317,6 +323,7 @@ func TestRemoteEvalErrorResponse(t *testing.T) {
317323

318324
ctx := context.Background()
319325
ctx = user.InjectOrgID(ctx, "test")
326+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
320327

321328
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
322329
require.ErrorContains(t, err, "rule evaluation failed")
@@ -343,6 +350,7 @@ func TestRemoteEvalNon2xxResponse(t *testing.T) {
343350

344351
ctx := context.Background()
345352
ctx = user.InjectOrgID(ctx, "test")
353+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
346354

347355
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
348356
require.ErrorContains(t, err, fmt.Sprintf("unsuccessful/unexpected response - status code %d", httpErr))
@@ -367,6 +375,7 @@ func TestRemoteEvalNonJSONResponse(t *testing.T) {
367375

368376
ctx := context.Background()
369377
ctx = user.InjectOrgID(ctx, "test")
378+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
370379

371380
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
372381
require.ErrorContains(t, err, "unexpected body encoding, not valid JSON")
@@ -406,6 +415,7 @@ func TestRemoteEvalUnsupportedResultResponse(t *testing.T) {
406415

407416
ctx := context.Background()
408417
ctx = user.InjectOrgID(ctx, "test")
418+
ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
409419

410420
_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
411421
require.ErrorContains(t, err, fmt.Sprintf("unsupported result type: %q", loghttp.ResultTypeStream))

0 commit comments

Comments
 (0)