feat: Report served-request usage - #19
gagantrivedi wants to merge 22 commits into
Conversation
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars. ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Advanced Run ID: 📝 WalkthroughWalkthroughThe proxy now validates and stores a usage flush interval, aggregates successful requests by client key and resource, and submits usage batches with retry and idempotency handling. Router state now contains environment and usage services. Middleware records supported successful requests. Environment document fetches include the proxy key when configured, and cross-origin pagination links are rejected. Unit and integration tests cover the new behaviour. Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🟡 Moderate · up to Proxy credentials can be exposed by accepted HTTP configurations or redirects, and concurrent configuration refreshes can produce incorrect usage records. These issues should be fixed 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.
| 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.
There was a problem hiding this comment.
Interesting, this could be a problem too. If the ingestion endpoint fails after persisting we could end up double counting.
There was a problem hiding this comment.
@matthewelwell, agreed. A timeout, connection failure, or 5xx response is ambiguous. The ingestion endpoint can persist the batch before the proxy observes the failure.
The proxy should retain the exact batch and resend it with the same stable idempotency key. The ingestion endpoint should persist that key atomically and treat later deliveries of that key as successful no-ops. The key must remain stable across retries and must not be derived only from the usage rows.
🐇🔍
You are interacting with an AI system.
| 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.
matthewelwell
left a comment
matthewelwell
left a comment
There was a problem hiding this comment.
Mostly just agreeing with the CR comments with one additional recommendation for desloppification.
| /// Report the counts accumulated since the last flush to the usage | ||
| /// endpoint, in chunks the server accepts. A rejected (4xx) chunk is | ||
| /// dropped — retrying cannot heal a rejection, and losing one window | ||
| /// beats resending a poisoned batch forever. Any other failure keeps | ||
| /// the chunk for the next flush. Returns false when any chunk was | ||
| /// not accepted. |
There was a problem hiding this comment.
I don't think this docstring really adds anything - the code describes all of this already
| 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.
Interesting, this could be a problem too. If the ingestion endpoint fails after persisting we could end up double counting.
90f27dd to
43c1d37
Compare
36d7cdd to
b64ca81
Compare
c3e23d1 to
dc037e7
Compare
The proxy counts every SDK request it serves, aggregated per environment
and resource inside resolve_key so no entry point can forget to count,
and flushes to POST {api_url}/proxy/usage/ every
usage_flush_interval_seconds (default 60) with the proxy key.
- flushes are chunked to the server's 1000-row cap; a rejected (4xx)
chunk is dropped, a failed (5xx/network) chunk is kept for next time
- document fetches carry X-Proxy-Key so core stops counting the proxy's
own polls
- statically configured environments keep their old billing: fetches
unmarked, served requests unreported
- inert without proxy_key; the final partial window is lost on shutdown
Counting inside resolve_key billed requests that then failed (404 on an unknown feature, 503 before the document loaded). A router middleware now counts after the handler answers, and only when it answered 2xx.
The next page request carries the environment and proxy keys, so a Link header must not be able to send them off api_url's origin.
A batch that failed with 5xx or a network error may still have been processed. It is now kept intact and resent with the same Idempotency-Key so the server can recognise it, instead of being merged into the next flush and counted twice. Nothing new is drained while a batch is pending; counts keep aggregating in the bounded map.
tokio intervals burst missed ticks by default, so a flush that overran the interval would be followed by back-to-back flushes.
c3e243a to
78d2f69
Compare
Static and discovered environments were billed differently: static ones kept unmarked fetches and unreported requests. With a proxy key, all fetches are now marked and all served requests reported; the server's grant check decides what counts.
3e1d7f2 swept criterion and a [[bench]] target into an unrelated commit; the bench file itself was never added.
EnvironmentService was also counting served requests and posting usage batches. That now lives in services/usage.rs; the router state is a struct holding both services, and the middleware resolves the key through one and records on the other.
|
@coderabbitai full review |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
🟠 Major · Require HTTPS when proxy_key is configured. · settings.rs:100-104
src/config/settings.rs:100-104
🔒 Security & Privacy | 🟠 Major | ⚡ Quick winRequire HTTPS when
proxy_keyis configured.
api_urlaccepts non-HTTPS URLs, while the usage, configuration, and document requests sendX-Proxy-Keyto that URL. Reject a non-HTTPSapi_urlwhenproxy_keyis set before starting credential-bearing reporting.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Advanced
Run ID: e8164de0-9861-474e-9ea4-fdd676e17ec5
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (21)
Cargo.tomlsrc/config/settings.rssrc/lib.rssrc/main.rssrc/middleware/cors.rssrc/middleware/mod.rssrc/middleware/usage.rssrc/routes/environment_document.rssrc/routes/flags.rssrc/routes/health.rssrc/routes/identities.rssrc/routes/mod.rssrc/services/environment.rssrc/services/mod.rssrc/services/usage.rssrc/state.rssrc/usage.rstests/test_cors.rstests/test_health.rstests/test_server.rstests/test_usage_tracking.rs
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.
|
reqwest strips only its own list of sensitive headers on a cross-host redirect; X-Proxy-Key and X-Environment-Key would be forwarded.
2 participants
Changes
Contributes to Flagsmith/flagsmith-private#256
The proxy counts every SDK request it serves with a 2xx — a router middleware counts after the handler answers, aggregated per environment and resource under the canonical client key, so failed requests (404, 503) and unresolved keys are never billed — and flushes to
POST {api_url}/proxy/usage/everyusage_flush_interval_seconds(default 60), authenticated by the proxy key.Idempotency-Key. A rejected (4xx) batch is dropped — retrying can't heal a rejection. A failed (5xx/network) batch may already have been processed, so it is kept intact and resent unchanged under the same key (Flagsmith/flagsmith-private#286 dedups it); nothing new is drained while a batch is pending, counts keep aggregating in the bounded map.X-Proxy-Keyso core stops counting the proxy's own polls. Arel=nextpagination link offapi_url's origin now fails the fetch instead of being followed, so neither key can be sent elsewhere.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.Companion PRs: usage ingestion endpoint (flagsmith-private) and the core middleware exclusion (links in the first comment).
How did you test this code?
85 tests (
cargo test), 10 new wiremock contract tests intests/test_usage_tracking.rs: aggregation across client/server keys through the full router; unresolved keys never counted; failed requests (404, 503) never counted; failed flush resends the same batch under the same key before new counts; rejected flush drops instead of retrying; 1001 environments chunk as 1000+1; static environment billed like a discovered one once a key is set; flush inert without a proxy key. clippy + fmt clean.