diff --git a/pkg/ruler/api.go b/pkg/ruler/api.go index fa4a1677ff9..76bde5f61d5 100644 --- a/pkg/ruler/api.go +++ b/pkg/ruler/api.go @@ -6,16 +6,17 @@ package ruler import ( + "encoding/hex" "encoding/json" "fmt" "io" "net/http" "net/url" - "sort" "strconv" "strings" "time" + "github.com/cespare/xxhash/v2" "github.com/go-kit/log" "github.com/go-kit/log/level" "github.com/gorilla/mux" @@ -64,6 +65,7 @@ type Alert struct { // RuleDiscovery has info for all rules type RuleDiscovery struct { RuleGroups []*RuleGroup `json:"groups"` + NextToken string `json:"nextToken,omitempty"` } // RuleGroup has info for rules which are part of a group @@ -166,6 +168,16 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) { return } + nextToken := req.URL.Query().Get("next_token") + var maxGroups int + if maxGroupsVal := req.URL.Query().Get("max_groups"); maxGroupsVal != "" { + maxGroups, err = strconv.Atoi(maxGroupsVal) + if err != nil || maxGroups < 0 { + respondInvalidRequest(logger, w, "invalid max groups value") + return + } + } + rulesReq := RulesRequest{ Filter: AnyRule, RuleName: req.URL.Query()["rule_name"], @@ -195,8 +207,21 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) { } groups := make([]*RuleGroup, 0, len(rgs)) - + var newToken string + foundToken := false for _, g := range rgs { + if nextToken != "" && !foundToken { + if nextToken != getRuleGroupNextToken(g.Group.Namespace, g.Group.Name) { + continue + } + foundToken = true + } + + if maxGroups > 0 && len(groups) == maxGroups { + newToken = getRuleGroupNextToken(g.Group.Namespace, g.Group.Name) + break + } + grp := RuleGroup{ Name: g.Group.Name, File: g.Group.Namespace, @@ -241,17 +266,13 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) { } } } + groups = append(groups, &grp) } - // keep data.groups are in order - sort.Slice(groups, func(i, j int) bool { - return groups[i].File < groups[j].File - }) - b, err := json.Marshal(&response{ Status: "success", - Data: &RuleDiscovery{RuleGroups: groups}, + Data: &RuleDiscovery{RuleGroups: groups, NextToken: newToken}, }) if err != nil { level.Error(logger).Log("msg", "error marshaling json response", "err", err) @@ -265,6 +286,13 @@ func (a *API) PrometheusRules(w http.ResponseWriter, req *http.Request) { } } +func getRuleGroupNextToken(file, group string) string { + h := xxhash.New() + _, _ = h.Write([]byte(file + ":" + group)) + + return hex.EncodeToString(h.Sum(nil)) +} + func (a *API) PrometheusAlerts(w http.ResponseWriter, req *http.Request) { logger, ctx := spanlogger.NewWithLogger(req.Context(), a.logger, "API.PrometheusAlerts") defer logger.Finish() diff --git a/pkg/ruler/api_test.go b/pkg/ruler/api_test.go index 10c63cbcf55..e04dd51beac 100644 --- a/pkg/ruler/api_test.go +++ b/pkg/ruler/api_test.go @@ -14,6 +14,7 @@ import ( "net/http" "net/http/httptest" "net/url" + "strconv" "strings" "testing" "time" @@ -905,6 +906,135 @@ func TestRuler_PrometheusRules(t *testing.T) { } } +func TestRuler_PrometheusRulesPagination(t *testing.T) { + const ( + userID = "user1" + interval = time.Minute + ) + + ruleGroups := rulespb.RuleGroupList{} + for ns := 0; ns < 3; ns++ { + for group := 0; group < 3; group++ { + g := &rulespb.RuleGroupDesc{ + Name: fmt.Sprintf("test-group-%d", group), + Namespace: fmt.Sprintf("test-namespace-%d", ns), + User: userID, + Rules: []*rulespb.RuleDesc{ + createAlertingRule("testalertingrule", "up < 1"), + }, + Interval: interval, + } + ruleGroups = append(ruleGroups, g) + } + } + + cfg := defaultRulerConfig(t) + cfg.TenantFederation.Enabled = true + + storageRules := map[string]rulespb.RuleGroupList{ + userID: ruleGroups, + } + + r := prepareRuler(t, cfg, newMockRuleStore(storageRules), withRulerAddrAutomaticMapping(), withLimits(validation.MockDefaultOverrides()), withStart()) + + // Rules will be synchronized asynchronously, so we wait until the expected number of rule groups + // has been synched. + test.Poll(t, 5*time.Second, len(ruleGroups), func() interface{} { + ctx := user.InjectOrgID(context.Background(), userID) + rls, _ := r.Rules(ctx, &RulesRequest{}) + return len(rls.Groups) + }) + + a := NewAPI(r, r.directStore, log.NewNopLogger()) + + getRulesResponse := func(groupSize int, nextToken string) response { + queryParams := "?" + url.Values{ + "max_groups": []string{strconv.Itoa(groupSize)}, + "next_token": []string{nextToken}, + }.Encode() + req := requestFor(t, http.MethodGet, "https://localhost:8080/prometheus/api/v1/rules"+queryParams, nil, userID) + w := httptest.NewRecorder() + a.PrometheusRules(w, req) + + resp := w.Result() + body, _ := io.ReadAll(resp.Body) + + r := response{} + err := json.Unmarshal(body, &r) + require.NoError(t, err) + + return r + } + + getRulesFromResponse := func(resp response) RuleDiscovery { + jsonRules, err := json.Marshal(resp.Data) + require.NoError(t, err) + returnedRules := RuleDiscovery{} + require.NoError(t, json.Unmarshal(jsonRules, &returnedRules)) + + return returnedRules + } + + // No page size limit + resp := getRulesResponse(0, "") + require.Equal(t, "success", resp.Status) + rd := getRulesFromResponse(resp) + require.Len(t, rd.RuleGroups, len(ruleGroups)) + require.Empty(t, rd.NextToken) + + // We have 9 groups, keep fetching rules with a group page size of 2. The final + // page should have size 1 and an empty nextToken. Also check the groups returned + // in order + var nextToken string + returnedRuleGroups := make([]*RuleGroup, 0, len(ruleGroups)) + for i := 0; i < 4; i++ { + resp := getRulesResponse(2, nextToken) + require.Equal(t, "success", resp.Status) + + rd := getRulesFromResponse(resp) + require.Len(t, rd.RuleGroups, 2) + require.NotEmpty(t, rd.NextToken) + + returnedRuleGroups = append(returnedRuleGroups, rd.RuleGroups[0], rd.RuleGroups[1]) + nextToken = rd.NextToken + } + resp = getRulesResponse(2, nextToken) + require.Equal(t, "success", resp.Status) + + rd = getRulesFromResponse(resp) + require.Len(t, rd.RuleGroups, 1) + require.Empty(t, rd.NextToken) + returnedRuleGroups = append(returnedRuleGroups, rd.RuleGroups[0]) + + // Check the returned rules match the rules written + require.Equal(t, len(ruleGroups), len(returnedRuleGroups)) + for i := 0; i < len(ruleGroups); i++ { + require.Equal(t, ruleGroups[i].Namespace, returnedRuleGroups[i].File) + require.Equal(t, ruleGroups[i].Name, returnedRuleGroups[i].Name) + require.Equal(t, len(ruleGroups[i].Rules), len(returnedRuleGroups[i].Rules)) + for j := 0; j < len(ruleGroups[i].Rules); j++ { + jsonRule, err := json.Marshal(returnedRuleGroups[i].Rules[j]) + require.NoError(t, err) + rule := alertingRule{} + require.NoError(t, json.Unmarshal(jsonRule, &rule)) + require.Equal(t, ruleGroups[i].Rules[j].Alert, rule.Name) + } + } + + // Invalid max groups value + resp = getRulesResponse(-1, "") + require.Equal(t, "error", resp.Status) + require.Equal(t, v1.ErrBadData, resp.ErrorType) + require.Equal(t, "invalid max groups value", resp.Error) + + // Bad token should return no groups + resp = getRulesResponse(0, "bad-token") + require.Equal(t, "success", resp.Status) + + rd = getRulesFromResponse(resp) + require.Len(t, rd.RuleGroups, 0) +} + func TestRuler_PrometheusAlerts(t *testing.T) { cfg := defaultRulerConfig(t)