Skip to content

Commit 2c55e43

Browse files
masterkainmasterkain
committed
perf(gateway): batch quota projection reads
1 parent 55bee4b commit 2c55e43

4 files changed

Lines changed: 368 additions & 11 deletions

File tree

‎lib/codex_pooler/alerts/evaluator.ex‎

Lines changed: 149 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@ defmodule CodexPooler.Alerts.Evaluator do
2828
}
2929

3030
@type evaluation_opts :: keyword() | map()
31+
@type projection_cache :: %{
32+
optional({Ecto.UUID.t() | nil, String.t() | nil}) => [map()]
33+
}
3134

3235
@spec evaluate_rule(AlertRule.t(), evaluation_opts()) :: [candidate()]
3336
def evaluate_rule(rule, opts \\ [])
@@ -95,15 +98,136 @@ defmodule CodexPooler.Alerts.Evaluator do
9598

9699
@spec evaluate_active_rules(evaluation_opts()) :: [candidate()]
97100
def evaluate_active_rules(opts \\ []) do
98-
AlertRule
99-
|> where([rule], rule.state == "active")
100-
|> order_by([rule], asc: rule.created_at, asc: rule.id)
101-
|> Repo.all()
102-
|> Enum.flat_map(&evaluate_rule(&1, opts))
101+
timestamp = evaluation_timestamp(opts)
102+
103+
{candidate_groups, _projection_cache} =
104+
AlertRule
105+
|> where([rule], rule.state == "active")
106+
|> order_by([rule], asc: rule.created_at, asc: rule.id)
107+
|> Repo.all()
108+
|> Enum.map_reduce(%{}, fn rule, projection_cache ->
109+
evaluate_rule_with_projection_cache(rule, timestamp, projection_cache)
110+
end)
111+
112+
List.flatten(candidate_groups)
113+
end
114+
115+
defp evaluate_rule_with_projection_cache(
116+
%AlertRule{state: "disabled"} = rule,
117+
timestamp,
118+
projection_cache
119+
) do
120+
{[clear_candidate(rule, dedupe_key_for_rule(rule, nil), timestamp)], projection_cache}
121+
end
122+
123+
defp evaluate_rule_with_projection_cache(
124+
%AlertRule{rule_kind: "pool_no_usable_assignments"} = rule,
125+
timestamp,
126+
projection_cache
127+
) do
128+
{projection, projection_cache} =
129+
pool_projection_from_cache(rule.pool_id, rule.model, timestamp, projection_cache)
130+
131+
dedupe_key = dedupe_key_for_rule(rule, nil)
132+
133+
candidates =
134+
if projection.usable_assignment_count == 0 do
135+
[pool_match_candidate(rule, dedupe_key, projection, "no_usable_assignments", timestamp)]
136+
else
137+
[clear_candidate(rule, dedupe_key, timestamp)]
138+
end
139+
140+
{candidates, projection_cache}
141+
end
142+
143+
defp evaluate_rule_with_projection_cache(
144+
%AlertRule{rule_kind: "pool_low_usable_assignments"} = rule,
145+
timestamp,
146+
projection_cache
147+
) do
148+
min_usable = rule.min_usable_assignments || 1
149+
150+
{projection, projection_cache} =
151+
pool_projection_from_cache(rule.pool_id, rule.model, timestamp, projection_cache)
152+
153+
dedupe_key = dedupe_key_for_rule(rule, nil)
154+
155+
candidates =
156+
if projection.usable_assignment_count > 0 and
157+
projection.usable_assignment_count < min_usable do
158+
[pool_match_candidate(rule, dedupe_key, projection, "low_usable_assignments", timestamp)]
159+
else
160+
[clear_candidate(rule, dedupe_key, timestamp)]
161+
end
162+
163+
{candidates, projection_cache}
164+
end
165+
166+
defp evaluate_rule_with_projection_cache(
167+
%AlertRule{rule_kind: "pool_all_assignments_in_state"} = rule,
168+
timestamp,
169+
projection_cache
170+
) do
171+
{projection, projection_cache} =
172+
pool_projection_from_cache(rule.pool_id, rule.model, timestamp, projection_cache)
173+
174+
dedupe_key = dedupe_key_for_rule(rule, nil)
175+
target_state = rule.target_state
176+
177+
candidates =
178+
if (target_state && projection.enabled_assignment_count > 0) and
179+
all_in_state?(projection, target_state) do
180+
[pool_match_candidate(rule, dedupe_key, projection, target_state, timestamp)]
181+
else
182+
[clear_candidate(rule, dedupe_key, timestamp)]
183+
end
184+
185+
{candidates, projection_cache}
186+
end
187+
188+
defp evaluate_rule_with_projection_cache(
189+
%AlertRule{rule_kind: rule_kind} = rule,
190+
timestamp,
191+
projection_cache
192+
)
193+
when rule_kind in ["upstream_quota_threshold", "upstream_auth_state"] do
194+
{assignments, projection_cache} =
195+
assigned_identity_projections_from_cache(
196+
rule.pool_id,
197+
rule.model,
198+
timestamp,
199+
projection_cache
200+
)
201+
202+
candidates =
203+
assignments
204+
|> Enum.reject(&(&1.assignment_status in @disabled_assignment_states))
205+
|> Enum.map(fn assignment ->
206+
case rule.rule_kind do
207+
"upstream_quota_threshold" -> threshold_candidate(rule, assignment, timestamp)
208+
"upstream_auth_state" -> auth_state_candidate(rule, assignment, timestamp)
209+
end
210+
end)
211+
212+
{candidates, projection_cache}
103213
end
104214

105215
defp pool_projection(pool_id, model, timestamp) do
106-
assignments = assigned_identity_projections(pool_id, model, timestamp)
216+
pool_projection_from_assignments(
217+
pool_id,
218+
model,
219+
assigned_identity_projections(pool_id, model, timestamp)
220+
)
221+
end
222+
223+
defp pool_projection_from_cache(pool_id, model, timestamp, projection_cache) do
224+
{assignments, projection_cache} =
225+
assigned_identity_projections_from_cache(pool_id, model, timestamp, projection_cache)
226+
227+
{pool_projection_from_assignments(pool_id, model, assignments), projection_cache}
228+
end
229+
230+
defp pool_projection_from_assignments(pool_id, model, assignments) do
107231
enabled = Enum.reject(assignments, &(&1.assignment_status in @disabled_assignment_states))
108232
usable = Enum.filter(enabled, & &1.usable_assignment?)
109233

@@ -118,6 +242,25 @@ defmodule CodexPooler.Alerts.Evaluator do
118242
}
119243
end
120244

245+
@spec assigned_identity_projections_from_cache(
246+
Ecto.UUID.t() | nil,
247+
String.t() | nil,
248+
DateTime.t(),
249+
projection_cache()
250+
) :: {[map()], projection_cache()}
251+
defp assigned_identity_projections_from_cache(pool_id, model, timestamp, projection_cache) do
252+
cache_key = {pool_id, model}
253+
254+
case Map.fetch(projection_cache, cache_key) do
255+
{:ok, assignments} ->
256+
{assignments, projection_cache}
257+
258+
:error ->
259+
assignments = assigned_identity_projections(pool_id, model, timestamp)
260+
{assignments, Map.put(projection_cache, cache_key, assignments)}
261+
end
262+
end
263+
121264
defp assigned_identity_projections(pool_id, model, timestamp) do
122265
assignments = assignment_rows(pool_id)
123266

‎lib/codex_pooler/gateway/routing/file_selection.ex‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ defmodule CodexPooler.Gateway.Routing.FileSelection do
1313
RoutingSelection
1414
}
1515

16+
alias CodexPooler.Gateway.Runtime.Dispatch.RouteState
1617
alias CodexPooler.Repo
1718
alias CodexPooler.Upstreams.Schemas.{PoolUpstreamAssignment, UpstreamIdentity}
1819

@@ -55,9 +56,22 @@ defmodule CodexPooler.Gateway.Routing.FileSelection do
5556
defp route_selection(auth, candidates, payload, request_options, endpoint) do
5657
model = model()
5758

59+
route_state =
60+
%{visible_model: model, candidates: candidates}
61+
|> RouteState.new()
62+
|> RouteState.preload_routing_snapshots(auth, model, request_options)
63+
5864
with {:ok, candidates} <- require_file_candidates(candidates, request_options),
59-
{:ok, candidates, request_options} <-
60-
filter_file_candidates(auth, model, candidates, payload, request_options, endpoint),
65+
{:ok, candidates, request_options, route_state} <-
66+
filter_file_candidates(
67+
auth,
68+
model,
69+
candidates,
70+
payload,
71+
request_options,
72+
endpoint,
73+
route_state
74+
),
6175
{:ok, selection} <-
6276
RoutingSelection.select_and_begin_circuit(%{
6377
auth: auth,
@@ -66,7 +80,8 @@ defmodule CodexPooler.Gateway.Routing.FileSelection do
6680
route_plan_input: RoutePlanInput.from_request_opts(request_options),
6781
endpoint: endpoint,
6882
payload: payload,
69-
request_options: request_options
83+
request_options: request_options,
84+
route_state: route_state
7085
}) do
7186
{:ok, selection}
7287
else
@@ -88,7 +103,15 @@ defmodule CodexPooler.Gateway.Routing.FileSelection do
88103

89104
defp require_file_candidates(candidates, _request_options), do: {:ok, candidates}
90105

91-
defp filter_file_candidates(auth, model, candidates, payload, request_options, endpoint) do
106+
defp filter_file_candidates(
107+
auth,
108+
model,
109+
candidates,
110+
payload,
111+
request_options,
112+
endpoint,
113+
route_state
114+
) do
92115
%{
93116
auth: auth,
94117
model: model,
@@ -98,7 +121,7 @@ defmodule CodexPooler.Gateway.Routing.FileSelection do
98121
candidates: candidates
99122
}
100123
|> CandidateEligibility.FilterInput.new()
101-
|> RouteFiltering.filter_candidates(quota_mode: :optional)
124+
|> RouteFiltering.filter_candidates(route_state, quota_mode: :optional)
102125
end
103126

104127
defp route_selection_error(%{status: status, code: code, message: message} = reason, opts) do

‎test/codex_pooler/alerts/evaluator_predicates_test.exs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ defmodule CodexPooler.Alerts.EvaluatorPredicatesTest do
88

99
alias CodexPooler.Alerts
1010
alias CodexPooler.Alerts.Schemas.AlertRule
11+
alias CodexPooler.Repo
1112
alias CodexPooler.Upstreams.Quota.Windows, as: QuotaWindows
1213

1314
test "quota target states stay separate for all-assignment predicates" do
@@ -114,6 +115,42 @@ defmodule CodexPooler.Alerts.EvaluatorPredicatesTest do
114115
assert [%{action: :clear}] = Alerts.evaluate_rule(usable_rule, at: timestamp)
115116
end
116117

118+
test "active rule evaluation reuses pool quota projections for the same pool" do
119+
timestamp = now()
120+
pool = pool_fixture()
121+
%{identity: identity} = upstream_assignment_fixture(pool)
122+
123+
assert {:ok, [_window]} =
124+
QuotaWindows.upsert_quota_windows(identity, [
125+
primary_quota_window_attrs(%{
126+
used_percent: Decimal.new("44"),
127+
credits: 56,
128+
reset_at: DateTime.add(timestamp, 1, :hour),
129+
observed_at: timestamp
130+
})
131+
])
132+
133+
alert_rule_fixture(pool, rule_kind: "pool_no_usable_assignments")
134+
135+
alert_rule_fixture(pool,
136+
rule_kind: "pool_low_usable_assignments",
137+
min_usable_assignments: 2
138+
)
139+
140+
alert_rule_fixture(pool,
141+
rule_kind: "pool_all_assignments_in_state",
142+
target_state: "missing_evidence"
143+
)
144+
145+
{_candidates, query_counts} =
146+
count_repo_commands(fn ->
147+
Alerts.evaluate_active_rules(at: timestamp)
148+
end)
149+
150+
assert command_count(query_counts, "pool_upstream_assignments", "SELECT") == 1
151+
assert command_count(query_counts, "account_quota_windows", "SELECT") == 1
152+
end
153+
117154
test "upstream quota threshold produces upstream-global metadata-only match candidates" do
118155
timestamp = now()
119156
pool = pool_fixture()
@@ -223,5 +260,51 @@ defmodule CodexPooler.Alerts.EvaluatorPredicatesTest do
223260
assert match.safe_evidence_snapshot["state_counts"][target_state] == 1
224261
end
225262

263+
defp count_repo_commands(fun) do
264+
parent = self()
265+
handler_id = "evaluator-predicates-test-#{System.unique_integer([:positive])}"
266+
267+
:ok =
268+
:telemetry.attach(
269+
handler_id,
270+
[:codex_pooler, :repo, :query],
271+
fn _event, _measurements, metadata, _config ->
272+
if metadata[:repo] == Repo do
273+
send(parent, {handler_id, metadata[:source], command_name(metadata[:query])})
274+
end
275+
end,
276+
nil
277+
)
278+
279+
try do
280+
result = fun.()
281+
{result, drain_repo_commands(handler_id, %{})}
282+
after
283+
:telemetry.detach(handler_id)
284+
end
285+
end
286+
287+
defp drain_repo_commands(handler_id, commands) do
288+
receive do
289+
{^handler_id, source, command} ->
290+
key = {source, command}
291+
drain_repo_commands(handler_id, Map.update(commands, key, 1, &(&1 + 1)))
292+
after
293+
0 -> commands
294+
end
295+
end
296+
297+
defp command_count(commands, source, command), do: Map.get(commands, {source, command}, 0)
298+
299+
defp command_name(query) when is_binary(query) do
300+
query
301+
|> String.trim_leading()
302+
|> String.split(~r/\s+/, parts: 2)
303+
|> List.first()
304+
|> String.upcase()
305+
end
306+
307+
defp command_name(_query), do: nil
308+
226309
defp now, do: DateTime.utc_now() |> DateTime.truncate(:microsecond)
227310
end

0 commit comments

Comments
 (0)