feat(transports): Category-based Rate Limiting (#354)

This adds support for parsing the X-Sentry-Rate-Limits and rate limiting errors and transactions independently.

internal/ratelimit/category.go

@@ -0,0 +1,41 @@
+package ratelimit
+
+import "strings"
+
+// Reference:
+// https://github.com/getsentry/relay/blob/0424a2e017d193a93918053c90cdae9472d164bf/relay-common/src/constants.rs#L116-L127
+
+// Category classifies supported payload types that can be ingested by Sentry
+// and, therefore, rate limited.
+type Category string
+
+// Known rate limit categories. As a special case, the CategoryAll applies to
+// all known payload types.
+const (
+	CategoryAll         Category = ""
+	CategoryError       Category = "error"
+	CategoryTransaction Category = "transaction"
+)
+
+// knownCategories is the set of currently known categories. Other categories
+// are ignored for the purpose of rate-limiting.
+var knownCategories = map[Category]struct{}{
+	CategoryAll:         {},
+	CategoryError:       {},
+	CategoryTransaction: {},
+}
+
+// String returns the category formatted for debugging.
+func (c Category) String() string {
+	switch c {
+	case "":
+		return "CategoryAll"
+	default:
+		var b strings.Builder
+		b.WriteString("Category")
+		for _, w := range strings.Fields(string(c)) {
+			b.WriteString(strings.Title(w))
+		}
+		return b.String()
+	}
+}

internal/ratelimit/category_test.go

@@ -0,0 +1,25 @@
+package ratelimit
+
+import "testing"
+
+func TestCategoryString(t *testing.T) {
+	tests := []struct {
+		Category
+		want string
+	}{
+		{CategoryAll, "CategoryAll"},
+		{CategoryError, "CategoryError"},
+		{CategoryTransaction, "CategoryTransaction"},
+		{Category("unknown"), "CategoryUnknown"},
+		{Category("two words"), "CategoryTwoWords"},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.want, func(t *testing.T) {
+			got := tt.Category.String()
+			if got != tt.want {
+				t.Errorf("got %q, want %q", got, tt.want)
+			}
+		})
+	}
+}

internal/ratelimit/deadline.go

@@ -0,0 +1,22 @@
+package ratelimit
+
+import "time"
+
+// A Deadline is a time instant when a rate limit expires.
+type Deadline time.Time
+
+// After reports whether the deadline d is after other.
+func (d Deadline) After(other Deadline) bool {
+	return time.Time(d).After(time.Time(other))
+}
+
+// Equal reports whether d and e represent the same deadline.
+func (d Deadline) Equal(e Deadline) bool {
+	return time.Time(d).Equal(time.Time(e))
+}
+
+// String returns the deadline formatted for debugging.
+func (d Deadline) String() string {
+	// Like time.Time.String, but without the monotonic clock reading.
+	return time.Time(d).Round(0).String()
+}

internal/ratelimit/doc.go

@@ -0,0 +1,3 @@
+// Package ratelimit provides tools to work with rate limits imposed by Sentry's
+// data ingestion pipeline.
+package ratelimit

internal/ratelimit/map.go

@@ -0,0 +1,64 @@
+package ratelimit
+
+import (
+	"net/http"
+	"time"
+)
+
+// Map maps categories to rate limit deadlines.
+//
+// A rate limit is in effect for a given category if either the category's
+// deadline or the deadline for the special CategoryAll has not yet expired.
+//
+// Use IsRateLimited to check whether a category is rate-limited.
+type Map map[Category]Deadline
+
+// IsRateLimited returns true if the category is currently rate limited.
+func (m Map) IsRateLimited(c Category) bool {
+	return m.isRateLimited(c, time.Now())
+}
+
+func (m Map) isRateLimited(c Category, now time.Time) bool {
+	return m.Deadline(c).After(Deadline(now))
+}
+
+// Deadline returns the deadline when the rate limit for the given category or
+// the special CategoryAll expire, whichever is furthest into the future.
+func (m Map) Deadline(c Category) Deadline {
+	categoryDeadline := m[c]
+	allDeadline := m[CategoryAll]
+	if categoryDeadline.After(allDeadline) {
+		return categoryDeadline
+	}
+	return allDeadline
+}
+
+// Merge merges the other map into m.
+//
+// If a category appears in both maps, the deadline that is furthest into the
+// future is preserved.
+func (m Map) Merge(other Map) {
+	for c, d := range other {
+		if d.After(m[c]) {
+			m[c] = d
+		}
+	}
+}
+
+// FromResponse returns a rate limit map from an HTTP response.
+func FromResponse(r *http.Response) Map {
+	return fromResponse(r, time.Now())
+}
+
+func fromResponse(r *http.Response, now time.Time) Map {
+	s := r.Header.Get("X-Sentry-Rate-Limits")
+	if s != "" {
+		return parseXSentryRateLimits(s, now)
+	}
+	if r.StatusCode == http.StatusTooManyRequests {
+		s := r.Header.Get("Retry-After")
+		deadline, _ := parseRetryAfter(s, now)
+		return Map{CategoryAll: deadline}
+	}
+	return Map{}
+}