From df3c44840120de6e269d759b0cdb69d6458f1ec5 Mon Sep 17 00:00:00 2001 From: jcalabro Date: Mon, 29 Jun 2026 12:00:35 -0400 Subject: [PATCH] Fix relay account limit alert claims --- cmd/relay/alerts.go | 26 +++++++++++++++++----- cmd/relay/alerts_test.go | 35 ++++++++++++++++++++++++++++++ cmd/relay/relay/host_usage.go | 24 ++++++++++++++------ cmd/relay/relay/host_usage_test.go | 20 +++++++++++++++++ 4 files changed, 92 insertions(+), 13 deletions(-) diff --git a/cmd/relay/alerts.go b/cmd/relay/alerts.go index ecb3375d..89366ea7 100644 --- a/cmd/relay/alerts.go +++ b/cmd/relay/alerts.go @@ -21,6 +21,7 @@ import ( ) const accountLimitAlertHostsPerMessage = 50 +const accountLimitAlertMaxTextChars = 2800 type AccountLimitAlertKind string @@ -257,12 +258,7 @@ func (m *AccountLimitAlertMonitor) Check(ctx context.Context) error { m.forwardAccountLimitAlertSent(ctx, accountLimitAlertSentState(AccountLimitAlertKindRecovery, usage, usage.AlertSentAt)) } - for start := 0; start < len(claims.Warnings); start += m.config.HostsPerMessage { - end := start + m.config.HostsPerMessage - if end > len(claims.Warnings) { - end = len(claims.Warnings) - } - batch := claims.Warnings[start:end] + for _, batch := range accountLimitAlertBatches(m.config, claims.Warnings) { if err := m.alerter.SendAlert(ctx, formatAccountLimitAlert(m.config, batch)); err != nil { return err } @@ -273,6 +269,24 @@ func (m *AccountLimitAlertMonitor) Check(ctx context.Context) error { return nil } +func accountLimitAlertBatches(config AccountLimitAlertMonitorConfig, usages []relay.HostAccountLimitUsage) [][]relay.HostAccountLimitUsage { + batches := make([][]relay.HostAccountLimitUsage, 0) + for _, usage := range usages { + if len(batches) == 0 { + batches = append(batches, []relay.HostAccountLimitUsage{usage}) + continue + } + last := batches[len(batches)-1] + candidate := append(append([]relay.HostAccountLimitUsage{}, last...), usage) + if len(last) >= config.HostsPerMessage || len([]rune(formatAccountLimitAlert(config, candidate).Text)) > accountLimitAlertMaxTextChars { + batches = append(batches, []relay.HostAccountLimitUsage{usage}) + continue + } + batches[len(batches)-1] = candidate + } + return batches +} + func (m *AccountLimitAlertMonitor) forwardAccountLimitAlertSent(ctx context.Context, state AccountLimitAlertSentState) { if m.config.SentCallback != nil { m.config.SentCallback(ctx, state) diff --git a/cmd/relay/alerts_test.go b/cmd/relay/alerts_test.go index 62fdbe45..9fc1413c 100644 --- a/cmd/relay/alerts_test.go +++ b/cmd/relay/alerts_test.go @@ -3,6 +3,7 @@ package main import ( "context" "encoding/json" + "fmt" "log/slog" "net/http" "net/http/httptest" @@ -189,6 +190,40 @@ func TestAccountLimitAlertMonitorBatchesMessagesWithoutQueryCap(t *testing.T) { assert.Contains(t, alerter.messages[1].Text, "two.example.com") } +func TestAccountLimitAlertMonitorBatchesMessagesBySlackTextLimit(t *testing.T) { + var warnings []relay.HostAccountLimitUsage + for i := 0; i < 80; i++ { + hostname := strings.Repeat("long-hostname-", 12) + fmt.Sprintf("%02d.example.com", i) + warnings = append(warnings, relay.HostAccountLimitUsage{ + ID: uint64(i + 1), + Hostname: hostname, + AccountCount: 90, + AccountLimit: 100, + Usage: 0.9, + AlertSentAt: time.Unix(1000, 0), + }) + } + claimer := &fakeAccountLimitAlertClaimer{ + claims: []relay.HostAccountLimitAlertClaims{ + {Warnings: warnings}, + }, + } + alerter := &recordingAlerter{} + config := DefaultAccountLimitAlertMonitorConfig() + + monitor, err := NewAccountLimitAlertMonitor(slog.Default(), claimer, alerter, config) + require.NoError(t, err) + require.NoError(t, monitor.Check(context.Background())) + + require.Greater(t, len(alerter.messages), 1) + seenHosts := 0 + for _, msg := range alerter.messages { + assert.LessOrEqual(t, len([]rune(msg.Text)), accountLimitAlertMaxTextChars) + seenHosts += strings.Count(msg.Text, ".example.com") + } + assert.Equal(t, len(warnings), seenHosts) +} + func TestSlackAlerterDoesNotPutTokenInPayload(t *testing.T) { const token = "xoxb-secret-token" var gotAuth string diff --git a/cmd/relay/relay/host_usage.go b/cmd/relay/relay/host_usage.go index 0d7ac19a..c5efc237 100644 --- a/cmd/relay/relay/host_usage.go +++ b/cmd/relay/relay/host_usage.go @@ -34,7 +34,7 @@ func (r *Relay) ListHostsApproachingAccountLimit(ctx context.Context, threshold var hosts []models.Host query := r.db.WithContext(ctx). Model(&models.Host{}). - Where("account_limit > 0 AND account_count >= account_limit * ? AND status <> ?", threshold, models.HostStatusBanned). + Where("account_limit > 0 AND account_count * 1.0 / account_limit >= ? AND status <> ?", threshold, models.HostStatusBanned). Order("account_count * 1.0 / account_limit DESC, id ASC") if limit > 0 { query = query.Limit(limit) @@ -96,7 +96,7 @@ func (r *Relay) claimDueAccountLimitWarnings(ctx context.Context, threshold floa var hosts []models.Host query := r.db.WithContext(ctx). Model(&models.Host{}). - Where("account_limit > 0 AND account_count >= account_limit * ? AND status <> ? AND account_limit_alerts_silenced = ?", threshold, models.HostStatusBanned, false). + Where("account_limit > 0 AND account_count * 1.0 / account_limit >= ? AND status <> ? AND account_limit_alerts_silenced = ?", threshold, models.HostStatusBanned, false). Where("(account_limit_alert_state <> ? OR account_limit_alert_state IS NULL OR account_limit_alert_sent_at IS NULL OR account_limit_alert_sent_at <= ?)", models.HostAccountLimitAlertStateWarning, cutoff). Order("account_count * 1.0 / account_limit DESC, id ASC") if limit > 0 { @@ -108,8 +108,13 @@ func (r *Relay) claimDueAccountLimitWarnings(ctx context.Context, threshold floa out := make([]HostAccountLimitUsage, 0, len(hosts)) for _, host := range hosts { + usage := hostAccountLimitUsage(host, now) + if usage.Usage < threshold { + continue + } + result := r.db.WithContext(ctx).Model(&models.Host{}). - Where("id = ? AND account_limit > 0 AND account_count >= account_limit * ? AND status <> ? AND account_limit_alerts_silenced = ?", host.ID, threshold, models.HostStatusBanned, false). + Where("id = ? AND account_limit > 0 AND account_count * 1.0 / account_limit >= ? AND status <> ? AND account_limit_alerts_silenced = ?", host.ID, threshold, models.HostStatusBanned, false). Where("(account_limit_alert_state <> ? OR account_limit_alert_state IS NULL OR account_limit_alert_sent_at IS NULL OR account_limit_alert_sent_at <= ?)", models.HostAccountLimitAlertStateWarning, cutoff). Updates(map[string]any{ "account_limit_alert_state": models.HostAccountLimitAlertStateWarning, @@ -121,7 +126,7 @@ func (r *Relay) claimDueAccountLimitWarnings(ctx context.Context, threshold floa if result.RowsAffected == 0 { continue } - out = append(out, hostAccountLimitUsage(host, now)) + out = append(out, usage) } return out, nil } @@ -130,7 +135,7 @@ func (r *Relay) claimDueAccountLimitRecoveries(ctx context.Context, threshold fl var hosts []models.Host query := r.db.WithContext(ctx). Model(&models.Host{}). - Where("account_limit > 0 AND account_count < account_limit * ? AND status <> ? AND account_limit_alerts_silenced = ? AND account_limit_alert_state = ?", threshold, models.HostStatusBanned, false, models.HostAccountLimitAlertStateWarning). + Where("account_limit > 0 AND account_count * 1.0 / account_limit < ? AND status <> ? AND account_limit_alerts_silenced = ? AND account_limit_alert_state = ?", threshold, models.HostStatusBanned, false, models.HostAccountLimitAlertStateWarning). Order("account_count * 1.0 / account_limit DESC, id ASC") if limit > 0 { query = query.Limit(limit) @@ -141,8 +146,13 @@ func (r *Relay) claimDueAccountLimitRecoveries(ctx context.Context, threshold fl out := make([]HostAccountLimitUsage, 0, len(hosts)) for _, host := range hosts { + usage := hostAccountLimitUsage(host, now) + if usage.Usage >= threshold { + continue + } + result := r.db.WithContext(ctx).Model(&models.Host{}). - Where("id = ? AND account_limit > 0 AND account_count < account_limit * ? AND status <> ? AND account_limit_alerts_silenced = ? AND account_limit_alert_state = ?", host.ID, threshold, models.HostStatusBanned, false, models.HostAccountLimitAlertStateWarning). + Where("id = ? AND account_limit > 0 AND account_count * 1.0 / account_limit < ? AND status <> ? AND account_limit_alerts_silenced = ? AND account_limit_alert_state = ?", host.ID, threshold, models.HostStatusBanned, false, models.HostAccountLimitAlertStateWarning). Updates(map[string]any{ "account_limit_alert_state": models.HostAccountLimitAlertStateOK, "account_limit_alert_sent_at": now, @@ -153,7 +163,7 @@ func (r *Relay) claimDueAccountLimitRecoveries(ctx context.Context, threshold fl if result.RowsAffected == 0 { continue } - out = append(out, hostAccountLimitUsage(host, now)) + out = append(out, usage) } return out, nil } diff --git a/cmd/relay/relay/host_usage_test.go b/cmd/relay/relay/host_usage_test.go index 1659f7df..5e8ae0d8 100644 --- a/cmd/relay/relay/host_usage_test.go +++ b/cmd/relay/relay/host_usage_test.go @@ -133,6 +133,26 @@ func TestClaimDueAccountLimitAlertsSkipsSilencedHosts(t *testing.T) { require.Empty(t, claims.Recoveries) } +func TestClaimDueAccountLimitAlertsDoesNotWarnAfterUnsilencingHostBelowThreshold(t *testing.T) { + ctx := context.Background() + r, db := testRelayWithHostDB(t) + host := models.Host{ + Hostname: "below-threshold.example.com", + AccountCount: 123667, + AccountLimit: 10000000, + Status: models.HostStatusActive, + AccountLimitAlertsSilenced: true, + } + require.NoError(t, db.Create(&host).Error) + + require.NoError(t, r.UpdateHostAccountLimitAlertsSilenced(ctx, host.ID, false)) + + claims, err := r.ClaimDueAccountLimitAlerts(ctx, 0.85, time.Hour, 0) + require.NoError(t, err) + require.Empty(t, claims.Warnings) + require.Empty(t, claims.Recoveries) +} + func TestUpdateHostAccountLimitAlertsSilencedResetsAlertState(t *testing.T) { ctx := context.Background() r, db := testRelayWithHostDB(t) -- 2.51.2