Skip to content

Commit d20ab45

Browse files
dmarticusgithub-actions[bot]
dmarticus
and
github-actions[bot]
authored
feat(flags): add quota limiting for feature flags served up with /decide and /local_evaluation (#28564)
Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
1 parent 8a1b1b0 commit d20ab45

File tree

11 files changed

+607
-141
lines changed

11 files changed

+607
-141
lines changed

‎ee/api/test/test_billing.py

Lines changed: 43 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from datetime import datetime
22
from typing import Any
3+
from unittest import TestCase
34
from unittest.mock import MagicMock, patch
45
from uuid import uuid4
56
from zoneinfo import ZoneInfo
@@ -50,6 +51,7 @@ def create_missing_billing_customer(**kwargs) -> CustomerInfo:
5051
"events": {"limit": None, "usage": 0},
5152
"recordings": {"limit": None, "usage": 0},
5253
"rows_synced": {"limit": None, "usage": 0},
54+
"feature_flag_requests": {"limit": None, "usage": 0},
5355
},
5456
free_trial_until=None,
5557
available_product_features=[],
@@ -145,6 +147,7 @@ def create_billing_customer(**kwargs) -> CustomerInfo:
145147
"events": {"limit": None, "usage": 0},
146148
"recordings": {"limit": None, "usage": 0},
147149
"rows_synced": {"limit": None, "usage": 0},
150+
"feature_flag_requests": {"limit": None, "usage": 0},
148151
},
149152
free_trial_until=None,
150153
)
@@ -436,6 +439,7 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
436439
"events": {"limit": None, "usage": 0},
437440
"recordings": {"limit": None, "usage": 0},
438441
"rows_synced": {"limit": None, "usage": 0},
442+
"feature_flag_requests": {"limit": None, "usage": 0},
439443
},
440444
"free_trial_until": None,
441445
}
@@ -559,6 +563,7 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
559563
"events": {"limit": None, "usage": 0},
560564
"recordings": {"limit": None, "usage": 0},
561565
"rows_synced": {"limit": None, "usage": 0},
566+
"feature_flag_requests": {"limit": None, "usage": 0},
562567
},
563568
"free_trial_until": None,
564569
"current_total_amount_usd": "0.00",
@@ -716,24 +721,32 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
716721
res = self.client.get("/api/billing")
717722
assert res.status_code == 200
718723
self.organization.refresh_from_db()
719-
assert self.organization.usage == {
720-
"events": {
721-
"limit": None,
722-
"todays_usage": 0,
723-
"usage": 1000,
724-
},
725-
"recordings": {
726-
"limit": None,
727-
"todays_usage": 0,
728-
"usage": 0,
729-
},
730-
"rows_synced": {
731-
"limit": None,
732-
"todays_usage": 0,
733-
"usage": 0,
724+
TestCase().assertDictEqual(
725+
self.organization.usage,
726+
{
727+
"events": {
728+
"limit": None,
729+
"todays_usage": 0,
730+
"usage": 1000,
731+
},
732+
"recordings": {
733+
"limit": None,
734+
"todays_usage": 0,
735+
"usage": 0,
736+
},
737+
"rows_synced": {
738+
"limit": None,
739+
"todays_usage": 0,
740+
"usage": 0,
741+
},
742+
"feature_flag_requests": {
743+
"limit": None,
744+
"todays_usage": 0,
745+
"usage": 0,
746+
},
747+
"period": ["2022-10-07T11:12:48", "2022-11-07T11:12:48"],
734748
},
735-
"period": ["2022-10-07T11:12:48", "2022-11-07T11:12:48"],
736-
}
749+
)
737750

738751
self.organization.usage = {"events": {"limit": None, "usage": 1000, "todays_usage": 1100000}}
739752
self.organization.save()
@@ -807,6 +820,7 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
807820
"events": {"limit": None, "usage": 0, "todays_usage": 0},
808821
"recordings": {"limit": None, "usage": 0, "todays_usage": 0},
809822
"rows_synced": {"limit": None, "usage": 0, "todays_usage": 0},
823+
"feature_flag_requests": {"limit": None, "usage": 0, "todays_usage": 0},
810824
"period": ["2022-10-07T11:12:48", "2022-11-07T11:12:48"],
811825
}
812826

@@ -831,14 +845,24 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
831845
mock_request.side_effect = mock_implementation
832846

833847
self.organization.customer_id = None
834-
self.organization.customer_trust_scores = {"recordings": 0, "events": 0, "rows_synced": 0}
848+
self.organization.customer_trust_scores = {
849+
"recordings": 0,
850+
"events": 0,
851+
"rows_synced": 0,
852+
"feature_flags": 0,
853+
}
835854
self.organization.save()
836855

837856
res = self.client.get("/api/billing")
838857
assert res.status_code == 200
839858
self.organization.refresh_from_db()
840859

841-
assert self.organization.customer_trust_scores == {"recordings": 0, "events": 15, "rows_synced": 0}
860+
assert self.organization.customer_trust_scores == {
861+
"recordings": 0,
862+
"events": 15,
863+
"rows_synced": 0,
864+
"feature_flags": 0,
865+
}
842866

843867
@patch("ee.api.billing.requests.get")
844868
def test_billing_with_supported_params(self, mock_get):

‎ee/billing/billing_manager.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,7 @@ def update_org_details(self, organization: Organization, billing_status: Billing
301301
events=usage_summary["events"],
302302
recordings=usage_summary["recordings"],
303303
rows_synced=usage_summary.get("rows_synced", {}),
304+
feature_flag_requests=usage_summary.get("feature_flag_requests", {}),
304305
period=[
305306
data["billing_period"]["current_period_start"],
306307
data["billing_period"]["current_period_end"],

‎ee/billing/quota_limiting.py

Lines changed: 45 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from posthog.exceptions_capture import capture_exception
1212

1313
from posthog.cache_utils import cache_for
14+
from posthog.constants import FlagRequestType
1415
from posthog.event_usage import report_organization_action
1516
from posthog.models.organization import Organization, OrganizationUsageInfo
1617
from posthog.models.team.team import Team
@@ -20,6 +21,7 @@
2021
get_teams_with_billable_event_count_in_period,
2122
get_teams_with_recording_count_in_period,
2223
get_teams_with_rows_synced_in_period,
24+
get_teams_with_feature_flag_requests_count_in_period,
2325
)
2426
from posthog.utils import get_current_day
2527

@@ -45,6 +47,7 @@ class QuotaResource(Enum):
4547
EVENTS = "events"
4648
RECORDINGS = "recordings"
4749
ROWS_SYNCED = "rows_synced"
50+
FEATURE_FLAG_REQUESTS = "feature_flag_requests"
4851

4952

5053
class QuotaLimitingCaches(Enum):
@@ -56,13 +59,22 @@ class QuotaLimitingCaches(Enum):
5659
QuotaResource.EVENTS: 0,
5760
QuotaResource.RECORDINGS: 1000,
5861
QuotaResource.ROWS_SYNCED: 0,
62+
QuotaResource.FEATURE_FLAG_REQUESTS: 0,
63+
}
64+
65+
TRUST_SCORE_KEYS = {
66+
QuotaResource.EVENTS: "events",
67+
QuotaResource.RECORDINGS: "recordings",
68+
QuotaResource.ROWS_SYNCED: "rows_synced",
69+
QuotaResource.FEATURE_FLAG_REQUESTS: "feature_flags",
5970
}
6071

6172

6273
class UsageCounters(TypedDict):
6374
events: int
6475
recordings: int
6576
rows_synced: int
77+
feature_flags: int
6678

6779

6880
# -------------------------------------------------------------------------------------------------
@@ -143,7 +155,9 @@ def org_quota_limited_until(
143155
quota_limiting_suspended_until = summary.get("quota_limiting_suspended_until", None)
144156
# Note: customer_trust_scores can initially be null. This should only happen after the initial migration and therefore
145157
# should be removed once all existing customers have this field set.
146-
trust_score = organization.customer_trust_scores.get(resource.value) if organization.customer_trust_scores else 0
158+
trust_score = (
159+
organization.customer_trust_scores.get(TRUST_SCORE_KEYS[resource]) if organization.customer_trust_scores else 0
160+
)
147161

148162
# Flow for checking quota limits:
149163
# 1. ignore the limits
@@ -367,7 +381,12 @@ def update_org_billing_quotas(organization: Organization):
367381
"update_org_billing_quotas started", {"today_end": today_end, "organization_id": organization.id}
368382
)
369383

370-
for resource in [QuotaResource.EVENTS, QuotaResource.RECORDINGS, QuotaResource.ROWS_SYNCED]:
384+
for resource in [
385+
QuotaResource.EVENTS,
386+
QuotaResource.RECORDINGS,
387+
QuotaResource.ROWS_SYNCED,
388+
QuotaResource.FEATURE_FLAG_REQUESTS,
389+
]:
371390
previously_quota_limited_team_tokens = list_limited_team_attributes(
372391
resource,
373392
QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY,
@@ -421,7 +440,7 @@ def set_org_usage_summary(
421440

422441
new_usage = copy.deepcopy(new_usage)
423442

424-
for field in ["events", "recordings", "rows_synced"]:
443+
for field in ["events", "recordings", "rows_synced", "feature_flag_requests"]:
425444
resource_usage = new_usage.get(field, {"limit": None, "usage": 0, "todays_usage": 0})
426445
if not resource_usage:
427446
continue
@@ -432,7 +451,7 @@ def set_org_usage_summary(
432451
org_usage_data = organization.usage or {}
433452
org_field_usage = org_usage_data.get(field, {}) or {}
434453
org_usage = org_field_usage.get("usage")
435-
# TRICKY: If we are not explictly setting todays_usage, we want to reset it to 0 IF the incoming new_usage is different
454+
# TRICKY: If we are not explicitly setting todays_usage, we want to reset it to 0 IF the incoming new_usage is different
436455
if org_usage != resource_usage.get("usage"):
437456
resource_usage["todays_usage"] = 0
438457
else:
@@ -476,6 +495,14 @@ def update_all_orgs_billing_quotas(
476495
"teams_with_rows_synced_in_period": convert_team_usage_rows_to_dict(
477496
get_teams_with_rows_synced_in_period(period_start, period_end)
478497
),
498+
"teams_with_decide_requests_count": convert_team_usage_rows_to_dict(
499+
get_teams_with_feature_flag_requests_count_in_period(period_start, period_end, FlagRequestType.DECIDE)
500+
),
501+
"teams_with_local_evaluation_requests_count": convert_team_usage_rows_to_dict(
502+
get_teams_with_feature_flag_requests_count_in_period(
503+
period_start, period_end, FlagRequestType.LOCAL_EVALUATION
504+
)
505+
),
479506
}
480507

481508
teams: Sequence[Team] = list(
@@ -503,10 +530,14 @@ def update_all_orgs_billing_quotas(
503530

504531
# we iterate through all teams, and add their usage to the organization they belong to
505532
for team in teams:
533+
decide_requests = all_data["teams_with_decide_requests_count"].get(team.id, 0)
534+
local_evaluation_requests = all_data["teams_with_local_evaluation_requests_count"].get(team.id, 0)
535+
506536
team_report = UsageCounters(
507537
events=all_data["teams_with_event_count_in_period"].get(team.id, 0),
508538
recordings=all_data["teams_with_recording_count_in_period"].get(team.id, 0),
509539
rows_synced=all_data["teams_with_rows_synced_in_period"].get(team.id, 0),
540+
feature_flags=decide_requests + (local_evaluation_requests * 10), # Same weighting as in _get_team_report
510541
)
511542

512543
org_id = str(team.organization.id)
@@ -521,7 +552,7 @@ def update_all_orgs_billing_quotas(
521552

522553
# Now we have the usage for all orgs for the current day
523554
# orgs_by_id is a dict of orgs by id (e.g. {"018e9acf-b488-0000-259c-534bcef40359": <Organization: 018e9acf-b488-0000-259c-534bcef40359>})
524-
# 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}})
555+
# 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_flag_requests": 100}})
525556
report_quota_limiting_event(
526557
"update_all_orgs_billing_quotas",
527558
{
@@ -544,14 +575,15 @@ def update_all_orgs_billing_quotas(
544575
QuotaResource(field), QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY
545576
)
546577
# We have the teams that are currently under quota limits
547-
# 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"]})
578+
# 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_flag_requests": ["phc_123", "phc_456"]})
548579
report_quota_limiting_event(
549580
"update_all_orgs_billing_quotas",
550581
{
551582
"event": "previously quota limited teams fetched",
552583
"events_count": len(previously_quota_limited_team_tokens["events"]),
553584
"recordings_count": len(previously_quota_limited_team_tokens["recordings"]),
554585
"rows_synced_count": len(previously_quota_limited_team_tokens["rows_synced"]),
586+
"feature_flags_count": len(previously_quota_limited_team_tokens["feature_flag_requests"]),
555587
},
556588
)
557589

@@ -564,7 +596,7 @@ def update_all_orgs_billing_quotas(
564596
if set_org_usage_summary(org, todays_usage=todays_report):
565597
org.save(update_fields=["usage"])
566598

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

578610
# Now we have the teams that are currently under quota limits
579-
# 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}})
580-
# 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}})
611+
# 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_flag_requests": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
612+
# 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_flag_requests": {"018e9acf-b488-0000-259c-534bcef40359": 1737867600}})
581613
report_quota_limiting_event(
582614
"update_all_orgs_billing_quotas",
583615
{
@@ -608,8 +640,8 @@ def update_all_orgs_billing_quotas(
608640
orgs_with_changes.add(org_id)
609641

610642
# Now we have the teams that are currently under quota limits
611-
# 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}})
612-
# 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}})
643+
# 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_flag_requests": {"phc_123": 1737867600}})
644+
# 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_flag_requests": {"phc_123": 1737867600}})
613645
report_quota_limiting_event(
614646
"update_all_orgs_billing_quotas",
615647
{
@@ -623,8 +655,9 @@ def update_all_orgs_billing_quotas(
623655
for org_id in orgs_with_changes:
624656
properties = {
625657
"quota_limited_events": quota_limited_orgs["events"].get(org_id, None),
626-
"quota_limited_recordings": quota_limited_orgs["events"].get(org_id, None),
658+
"quota_limited_recordings": quota_limited_orgs["recordings"].get(org_id, None),
627659
"quota_limited_rows_synced": quota_limited_orgs["rows_synced"].get(org_id, None),
660+
"quota_limited_feature_flags": quota_limited_orgs["feature_flag_requests"].get(org_id, None),
628661
}
629662

630663
report_organization_action(

0 commit comments

Comments
 (0)