Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(flags): add quota limiting for feature flags served up with /decide and to local_evaluation #28564

Open
wants to merge 28 commits into
base: master
Choose a base branch
from
+546 −126
Open

feat(flags): add quota limiting for feature flags served up with /decide and to local_evaluation #28564

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
e5b00e8
ruff
dmarticus Feb 11, 2025
03f877c
import constants
dmarticus Feb 11, 2025
20d7c14
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 11, 2025
2ed2be7
add feature flags to these quota limit assertions
dmarticus Feb 11, 2025
0dd8563
Merge branch 'feat/quota-limiting' of github.com:PostHog/posthog into…
dmarticus Feb 11, 2025
51f0df5
darnit i missed one
dmarticus Feb 11, 2025
b428fa8
two more bugs I fixed wow go me
dmarticus Feb 11, 2025
290d681
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 11, 2025
727c4cf
fixing some tests
dmarticus Feb 11, 2025
3836ffe
Merge branch 'feat/quota-limiting' of github.com:PostHog/posthog into…
dmarticus Feb 11, 2025
62b09c1
sigh idk
dmarticus Feb 11, 2025
b9c38d1
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 12, 2025
fef5963
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 12, 2025
ca1037b
Merge branch 'feat/quota-limiting' of github.com:PostHog/posthog into…
dmarticus Feb 12, 2025
e980e11
cleaning up quota limiting + remote config
dmarticus Feb 12, 2025
cd37a8b
more response polishing
dmarticus Feb 12, 2025
1d662c9
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 12, 2025
7d82a58
delete prints
dmarticus Feb 12, 2025
96f59ae
moar tests
dmarticus Feb 12, 2025
efd76bc
Merge branch 'master' into feat/quota-limiting
dmarticus Feb 12, 2025
9f06adf
limit local eval, too
dmarticus Feb 13, 2025
12c13c8
Merge branch 'feat/quota-limiting' of github.com:PostHog/posthog into…
dmarticus Feb 13, 2025
a176607
renamed the quota limit to make it more clear
dmarticus Feb 14, 2025
f1eab3f
respond with an error instead; this feels more honest
dmarticus Feb 14, 2025
641988e
resolvin them conflicts
dmarticus Feb 14, 2025
1f98d8b
Update query snapshots
github-actions[bot] Feb 14, 2025
14e4312
resolve conflicts
dmarticus Feb 14, 2025
d783b42
Update query snapshots
github-actions[bot] Feb 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 34 additions & 10 deletions ee/billing/quota_limiting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from posthog.exceptions_capture import capture_exception

from posthog.cache_utils import cache_for
from posthog.constants import FlagRequestType
from posthog.event_usage import report_organization_action
from posthog.models.organization import Organization, OrganizationUsageInfo
from posthog.models.team.team import Team
Expand All @@ -20,6 +21,7 @@
get_teams_with_billable_event_count_in_period,
get_teams_with_recording_count_in_period,
get_teams_with_rows_synced_in_period,
get_teams_with_feature_flag_requests_count_in_period,
)
from posthog.utils import get_current_day

Expand All @@ -45,6 +47,7 @@ class QuotaResource(Enum):
EVENTS = "events"
RECORDINGS = "recordings"
ROWS_SYNCED = "rows_synced"
FEATURE_FLAGS = "feature_flags"


class QuotaLimitingCaches(Enum):
Expand All @@ -56,13 +59,15 @@ class QuotaLimitingCaches(Enum):
QuotaResource.EVENTS: 0,
QuotaResource.RECORDINGS: 1000,
QuotaResource.ROWS_SYNCED: 0,
QuotaResource.FEATURE_FLAGS: 0, # TODO: should we have a buffer here?
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm open to calling this feature_flags_evaluated instead, that might be more clear (I like how rows_synced is more obvious than just events.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think FEATURE_FLAGS_EVALUATED makes sense. That way it's not confused with how many feature flags a team has, which I assume is NOT quota limited?

}


class UsageCounters(TypedDict):
events: int
recordings: int
rows_synced: int
feature_flags: int


# -------------------------------------------------------------------------------------------------
Expand Down Expand Up @@ -367,7 +372,12 @@ def update_org_billing_quotas(organization: Organization):
"update_org_billing_quotas started", {"today_end": today_end, "organization_id": organization.id}
)

for resource in [QuotaResource.EVENTS, QuotaResource.RECORDINGS, QuotaResource.ROWS_SYNCED]:
for resource in [
QuotaResource.EVENTS,
QuotaResource.RECORDINGS,
QuotaResource.ROWS_SYNCED,
QuotaResource.FEATURE_FLAGS,
]:
previously_quota_limited_team_tokens = list_limited_team_attributes(
resource,
QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY,
Expand Down Expand Up @@ -421,7 +431,7 @@ def set_org_usage_summary(

new_usage = copy.deepcopy(new_usage)

for field in ["events", "recordings", "rows_synced"]:
for field in ["events", "recordings", "rows_synced", "feature_flags"]:
resource_usage = new_usage.get(field, {"limit": None, "usage": 0, "todays_usage": 0})
if not resource_usage:
continue
Expand Down Expand Up @@ -476,6 +486,14 @@ def update_all_orgs_billing_quotas(
"teams_with_rows_synced_in_period": convert_team_usage_rows_to_dict(
get_teams_with_rows_synced_in_period(period_start, period_end)
),
"teams_with_decide_requests_count": convert_team_usage_rows_to_dict(
get_teams_with_feature_flag_requests_count_in_period(period_start, period_end, FlagRequestType.DECIDE)
),
"teams_with_local_evaluation_requests_count": convert_team_usage_rows_to_dict(
get_teams_with_feature_flag_requests_count_in_period(
period_start, period_end, FlagRequestType.LOCAL_EVALUATION
)
),
}

teams: Sequence[Team] = list(
Expand Down Expand Up @@ -503,10 +521,14 @@ def update_all_orgs_billing_quotas(

# we iterate through all teams, and add their usage to the organization they belong to
for team in teams:
decide_requests = all_data["teams_with_decide_requests_count"].get(team.id, 0)
local_evaluation_requests = all_data["teams_with_local_evaluation_requests_count"].get(team.id, 0)

team_report = UsageCounters(
events=all_data["teams_with_event_count_in_period"].get(team.id, 0),
recordings=all_data["teams_with_recording_count_in_period"].get(team.id, 0),
rows_synced=all_data["teams_with_rows_synced_in_period"].get(team.id, 0),
feature_flags=decide_requests + (local_evaluation_requests * 10), # Same weighting as in _get_team_report
)

org_id = str(team.organization.id)
Expand All @@ -521,7 +543,7 @@ def update_all_orgs_billing_quotas(

# Now we have the usage for all orgs for the current day
# orgs_by_id is a dict of orgs by id (e.g. {"018e9acf-b488-0000-259c-534bcef40359": <Organization: 018e9acf-b488-0000-259c-534bcef40359>})
# todays_usage_report is a dict of orgs by id with their usage for the current day (e.g. {"018e9acf-b488-0000-259c-534bcef40359": {"events": 100, "recordings": 100, "rows_synced": 100}})
# todays_usage_report is a dict of orgs by id with their usage for the current day (e.g. {"018e9acf-b488-0000-259c-534bcef40359": {"events": 100, "recordings": 100, "rows_synced": 100, "feature_flags": 100}})
report_quota_limiting_event(
"update_all_orgs_billing_quotas",
{
Expand All @@ -544,14 +566,15 @@ def update_all_orgs_billing_quotas(
QuotaResource(field), QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY
)
# We have the teams that are currently under quota limits
# previously_quota_limited_team_tokens is a dict of resources to team tokens from redis (e.g. {"events": ["phc_123", "phc_456"], "recordings": ["phc_123", "phc_456"], "rows_synced": ["phc_123", "phc_456"]})
# previously_quota_limited_team_tokens is a dict of resources to team tokens from redis (e.g. {"events": ["phc_123", "phc_456"], "recordings": ["phc_123", "phc_456"], "rows_synced": ["phc_123", "phc_456"], "feature_flags": ["phc_123", "phc_456"]})
report_quota_limiting_event(
"update_all_orgs_billing_quotas",
{
"event": "previously quota limited teams fetched",
"events_count": len(previously_quota_limited_team_tokens["events"]),
"recordings_count": len(previously_quota_limited_team_tokens["recordings"]),
"rows_synced_count": len(previously_quota_limited_team_tokens["rows_synced"]),
"feature_flags_count": len(previously_quota_limited_team_tokens["feature_flags"]),
},
)

Expand All @@ -564,7 +587,7 @@ def update_all_orgs_billing_quotas(
if set_org_usage_summary(org, todays_usage=todays_report):
org.save(update_fields=["usage"])

for field in ["events", "recordings", "rows_synced"]:
for field in ["events", "recordings", "rows_synced", "feature_flags"]:
# for each organization, we check if the current usage + today's unreported usage is over the limit
result = org_quota_limited_until(org, QuotaResource(field), previously_quota_limited_team_tokens[field])
if result:
Expand All @@ -576,8 +599,8 @@ def update_all_orgs_billing_quotas(
quota_limited_orgs[field][org_id] = quota_limited_until

# Now we have the teams that are currently under quota limits
# quota_limited_orgs is a dict of resources to org ids (e.g. {"events": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "recordings": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "rows_synced": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
# quota_limiting_suspended_orgs is a dict of resources to org ids (e.g. {"events": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "recordings": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "rows_synced": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
# quota_limited_orgs is a dict of resources to org ids (e.g. {"events": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "recordings": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "rows_synced": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "feature_flags": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
# quota_limiting_suspended_orgs is a dict of resources to org ids (e.g. {"events": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "recordings": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "rows_synced": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}, "feature_flags": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
report_quota_limiting_event(
"update_all_orgs_billing_quotas",
{
Expand Down Expand Up @@ -608,8 +631,8 @@ def update_all_orgs_billing_quotas(
orgs_with_changes.add(org_id)

# Now we have the teams that are currently under quota limits
# quota_limited_teams is a dict of resources to team tokens (e.g. {"events": {"phc_123": 1737867600}, "recordings": {"phc_123": 1737867600}, "rows_synced": {"phc_123": 1737867600}})
# quota_limiting_suspended_teams is a dict of resources to team tokens (e.g. {"events": {"phc_123": 1737867600}, "recordings": {"phc_123": 1737867600}, "rows_synced": {"phc_123": 1737867600}})
# quota_limited_teams is a dict of resources to team tokens (e.g. {"events": {"phc_123": 1737867600}, "recordings": {"phc_123": 1737867600}, "rows_synced": {"phc_123": 1737867600}, "feature_flags": {"phc_123": 1737867600}})
# quota_limiting_suspended_teams is a dict of resources to team tokens (e.g. {"events": {"phc_123": 1737867600}, "recordings": {"phc_123": 1737867600}, "rows_synced": {"phc_123": 1737867600}, "feature_flags": {"phc_123": 1737867600}})
report_quota_limiting_event(
"update_all_orgs_billing_quotas",
{
Expand All @@ -623,8 +646,9 @@ def update_all_orgs_billing_quotas(
for org_id in orgs_with_changes:
properties = {
"quota_limited_events": quota_limited_orgs["events"].get(org_id, None),
"quota_limited_recordings": quota_limited_orgs["events"].get(org_id, None),
"quota_limited_recordings": quota_limited_orgs["recordings"].get(org_id, None),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a bug fix? Has recordings quota limited been wrong this whole time?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh wow I didn't even realize i did this; it was a cursor tab-complete for adding in the quota_limited_flags stuff. I'm gonna tag the replay team so they know about this, since i don't want to change it without them signing off, but it's also possible they weren't even aware this was broken. @PostHog/team-replay you folks want to take a look at this?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🙈 that wouldn't suprise me.

i've never touched quota limiting and don't know how I'd check if it was broken 🤔

more scary, there were no failing tests before or after the change 😱

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

turns out it's been that way since being introduced

#14340

that PR talks about tracking so is this just what we're reporting to ourselves or the actual value we use for limiting?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol I wonder if it's because the one test where we check the value of that entry we also assert that it's the same value of the quota_limited_events field.

assert org_action_call.kwargs.get("properties") == {
    "quota_limited_events": 1612137599,
    "quota_limited_recordings": 1612137599,
    "quota_limited_rows_synced": None,
}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🙈

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you nailed it though; that object is just used in report_organization_action, so it's just for internal reporting

    for org_id in orgs_with_changes:
        properties = {
            "quota_limited_events": quota_limited_orgs["events"].get(org_id, None),
            "quota_limited_recordings": quota_limited_orgs["recordings"].get(org_id, None),
            "quota_limited_rows_synced": quota_limited_orgs["rows_synced"].get(org_id, None),
            "quota_limited_feature_flags": quota_limited_orgs["feature_flags"].get(org_id, None),
        }

        report_organization_action(
            orgs_by_id[org_id],
            "organization quota limits changed",
            properties=properties,
            group_properties=properties,
        )

So, no actual quota limiting bug, but a quota limit reporting bug!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are you ok to fix the test as well here?
i don't understand the impact but i'm guessing quota limiting on the correct count is better than not regardless of if it's just reporting

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh yeah i'm on it; there's a bunch of other tests I want to fix too – I need to assert the presence of this new quota limit in way more places lol

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

awesome... thanks!

"quota_limited_rows_synced": quota_limited_orgs["rows_synced"].get(org_id, None),
"quota_limited_feature_flags": quota_limited_orgs["feature_flags"].get(org_id, None),
}

report_organization_action(
Expand Down
17 changes: 17 additions & 0 deletions ee/billing/test/test_quota_limiting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -814,3 +814,20 @@ def test_get_team_attribute_by_quota_resource(self, mock_capture):
f"quota_limiting: No team tokens found for organization: {self.organization.id}",
str(mock_capture.call_args[0][0]),
)

def test_feature_flags_quota_limiting(self):
with self.settings(USE_TZ=False), freeze_time("2021-01-25T00:00:00Z"):
self.organization.usage = {
"events": {"usage": 10, "limit": 100},
"recordings": {"usage": 10, "limit": 100},
"feature_flags": {"usage": 110, "limit": 100},
"period": ["2021-01-01T00:00:00Z", "2021-01-31T23:59:59Z"],
}
self.organization.customer_trust_scores = {"events": 0, "recordings": 0, "feature_flags": 0}
self.organization.save()

quota_limited_orgs, quota_limiting_suspended_orgs = update_all_orgs_billing_quotas()

assert self.team.api_token.encode("UTF-8") in self.redis_client.zrange(
f"@posthog/quota-limits/feature_flags", 0, -1
)
dmarticus marked this conversation as resolved.
Show resolved Hide resolved
Loading
Loading