feat: Report served-request usage - #19
Conversation
Counts every request whose environment key resolves, aggregated per environment and resource, and flushes them to POST /proxy/usage/ every usage_flush_interval_seconds (default 60). Document fetches now carry X-Proxy-Key so core can exclude the proxy's own polls from API usage. Active only when proxy_key is configured; static config-file mode is unchanged. A final partial window is lost on shutdown (no graceful shutdown hook yet) — usage metering tolerates that.
A failed flush used to retry the identical batch forever, so any permanent rejection — a batch over the server's 1000-row cap, a revoked proxy key — silently ended usage reporting until restart. Flushes are now chunked to the server's cap and a rejected (4xx) chunk is dropped; only network errors and 5xx keep rows for the next flush. Also per review: a statically configured environment keeps its old billing end to end (its document fetch is unmarked and its served requests are not reported); counting moved into resolve_key(resource) so an SDK entry point cannot forget it; UsageCount renamed UsageRow; poll_usage renamed flush_usage_periodically; a zero flush interval now fails config validation instead of silently killing the task.
📝 WalkthroughWalkthroughThe change adds configurable usage flush intervals and public usage types. Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: 🟡 Moderate · up to Usage totals can be inaccurate, retry behavior can pressure ingestion, and a cross-origin pagination response can receive the proxy credential. These material issues should be resolved before merge. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Companion PRs for Flagsmith/flagsmith-private#256:
Merge/deploy order: #282 and 8407 (with the flagsmith-private pin bump) before any proxy build with #19 is deployed, else document polls double-count during the overlap. |
|
@coderabbitai review |
|
|
|
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Team
Run ID: beca0e8d-d764-4d52-aa88-31c6fd00bca0
📒 Files selected for processing (7)
src/config/settings.rssrc/environments.rssrc/lib.rssrc/main.rssrc/services/environment.rssrc/usage.rstests/test_usage_tracking.rs
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.
| .resolve(environment_key) | ||
| .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) | ||
| .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; | ||
| self.track_usage(&keys.client_key, resource); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Count usage only after a successful endpoint response.
resolve_key increments usage before get_flags_response_data, get_identity_response_data, or get_environment_bytes complete. A resolved key can still return ServiceUnavailable, and the flags path can return FeatureNotFound. These failed responses are counted as served requests. Move tracking to each successful Ok path, including cache hits.
| // unmarked and keep their old billing. | ||
| if !self.environments.is_static(server_side_key) { | ||
| if let Some(proxy_key) = &self.settings.proxy_key { | ||
| request = request.header("X-Proxy-Key", proxy_key); |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
sed -n '200,315p' src/services/environment.rs
printf '\n-- pagination symbols --\n'
rg -n -C 4 'parse_next_link|next_url|Link|X-Proxy-Key|fetch_environment' src/services/environment.rsRepository: Flagsmith/edge-proxy-rs
Length of output: 10695
🏁 Script executed:
#!/bin/bash
set -eu
sed -n '703,760p' src/services/environment.rs
printf '\n-- client dependency declarations --\n'
rg -n -C 3 'reqwest|^tokio|redirect' Cargo.toml Cargo.lockRepository: Flagsmith/edge-proxy-rs
Length of output: 3584
Sensitive Data Exposure (CWE-200): Exposure of Sensitive Information to an Unauthorized Actor
Reachability: External · Exploitability: Difficult
Restrict pagination URLs before sending X-Proxy-Key.
Accept only same-origin pagination URLs before forwarding the proxy credential. An absolute rel=next URL from the upstream Link header can otherwise direct the credential-bearing request to another origin.
| self.usage.merge(chunk); | ||
| all_success = false; | ||
| } | ||
| Err(e) => { | ||
| error!("Failed to report usage: {}", e); | ||
| self.usage.merge(chunk); | ||
| all_success = false; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Make retried usage batches idempotent.
Lines 634 and 639 retry a batch after a 5xx response or network error. Neither outcome proves that the ingestion endpoint did not persist the batch. UsageRow has no delivery identifier, and merge combines the failed rows with later usage. A retry can therefore over-count usage. Retain the exact batch with a stable idempotency key that the ingestion endpoint recognises before retrying it.
| let mut interval = tokio::time::interval(Duration::from_secs( | ||
| self.settings.usage_flush_interval_seconds, | ||
| )); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -A2 '^name = "tokio"$' Cargo.lockRepository: Flagsmith/edge-proxy-rs
Length of output: 275
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '600,690p' src/services/environment.rs
rg -n "flush_usage|MissedTickBehavior|interval\\(" src Cargo.toml Cargo.lockRepository: Flagsmith/edge-proxy-rs
Length of output: 4058
🌐 Web query:
Tokio 1.48.0 interval missed tick behavior default Burst set_missed_tick_behavior official docs
💡 Result:
In Tokio 1.48.0, the default behavior for an Interval when a tick is missed is MissedTickBehavior::Burst [1][2]. When the Burst strategy is active, the Interval will fire ticks as quickly as possible until it has "caught up" to the expected schedule [1][3]. This ensures that the ticks produced correspond to the same instants that would have been yielded if no ticks had been missed [1][3]. You can configure this behavior using the set_missed_tick_behavior method on an Interval instance [4][5]: interval.set_missed_tick_behavior(MissedTickBehavior::Skip); The MissedTickBehavior enum provides three strategies: - Burst (Default): Fires ticks as fast as possible to catch up [1][3]. - Skip: Skips missed ticks and schedules the next tick at the next multiple of the period [3]. - Delay: Reschedules all future ticks relative to the time the delayed tick was actually called [3]. The Burst strategy is the default to maintain backwards compatibility and because it is the desired behavior for most use cases [1][6]. Note that these strategies apply specifically when the delay exceeds 5 milliseconds due to executor precision limitations [1].
Citations:
- 1: https://docs.rs/tokio/latest/tokio/time/enum.MissedTickBehavior.html
- 2: https://docs.rs/tokio/latest/src/tokio/time/interval.rs.html
- 3: https://docs.rs/tokio/latest/%20tokio/time/enum.MissedTickBehavior.html
- 4: https://docs.rs/tokio/latest/tokio/time/struct.Interval.html
- 5: https://docs.rs/tokio/latest/tokio/time/fn.interval.html
- 6: https://github.com/tokio-rs/tokio/blob/c637f6e7/tokio/src/time/interval.rs
Use a non-burst missed-tick policy.
Tokio 1.48.0 uses MissedTickBehavior::Burst by default. If flush_usage exceeds the configured interval, subsequent tick() calls can complete immediately and retry failed reports without the configured delay. Set MissedTickBehavior::Delay or Skip.
2 participants
Changes
Contributes to Flagsmith/flagsmith-private#256
The proxy counts every SDK request it serves — aggregated per environment and resource, keyed by the canonical client key inside
resolve_key, so an entry point can't forget to count and unresolved keys never grow the map — and flushes toPOST {api_url}/proxy/usage/everyusage_flush_interval_seconds(default 60), authenticated by the proxy key.X-Proxy-Keyso core stops counting the proxy's own polls.proxy_keyis set; static config-file mode is byte-identical. The final partial window is lost on shutdown (no graceful-shutdown hook) — usage metering tolerates that.Stacked on #18. Companion PRs: usage ingestion endpoint (flagsmith-private) and the core middleware exclusion (links in the first comment).
How did you test this code?
92 tests (
cargo test), 9 new wiremock contract tests intests/test_usage_tracking.rs: aggregation across client/server keys through the full router; unresolved keys never counted; failed flush merges into the next (nothing lost, nothing doubled); rejected flush drops instead of retrying; 1001 environments chunk as 1000+1; static environment neither counted nor marked; flush inert without a proxy key. clippy + fmt clean.