Skip to content

Commit

Permalink
support optional concurrency policies
Browse files Browse the repository at this point in the history
  • Loading branch information
glendc committed Mar 9, 2024
1 parent d48ac2f commit f04fb97
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 2 deletions.
10 changes: 8 additions & 2 deletions examples/http_rate_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,14 +66,20 @@ async fn main() {
// but you can make them also optional to not use backoff for some, while using it for others
.add(
HttpMatcher::socket(SocketMatcher::loopback()).negate(),
ConcurrentPolicy::with_backoff(1, None),
Some(ConcurrentPolicy::with_backoff(1, None)),
)
// you can also use options for the policy itself, in case you want to disable
// the limit for some
.add(HttpMatcher::path("/admin/*"), None)
// test path so you can test also rate limiting on an http level
// > NOTE: as you can also make your own Matchers you can limit on w/e
// > property you want.
.add(
HttpMatcher::path("/limit/*"),
ConcurrentPolicy::with_backoff(2, Some(ExponentialBackoff::default())),
Some(ConcurrentPolicy::with_backoff(
2,
Some(ExponentialBackoff::default()),
)),
)
.build(),
))
Expand Down
47 changes: 47 additions & 0 deletions src/service/layer/limit/policy/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
//! Limit policies for [`super::Limit`]
//! define how requests are handled when the limit is reached
//! for a given request.
//!
//! [`Option`] can be used to disable a limit policy for some scenarios
//! while enabling it for others.

use crate::service::Context;

Expand Down Expand Up @@ -62,3 +65,47 @@ pub trait Policy<State, Request>: Send + Sync + 'static {
+ Send
+ '_;
}

impl<State, Request, P> Policy<State, Request> for Option<P>
where
P: Policy<State, Request>,
State: Send + Sync + 'static,
Request: Send + 'static,
{
type Guard = Option<P::Guard>;
type Error = P::Error;

async fn check(
&self,
ctx: Context<State>,
request: Request,
) -> PolicyResult<State, Request, Self::Guard, Self::Error> {
match self {
Some(policy) => {
let result = policy.check(ctx, request).await;
match result.output {
PolicyOutput::Ready(guard) => PolicyResult {
ctx: result.ctx,
request: result.request,
output: PolicyOutput::Ready(Some(guard)),
},
PolicyOutput::Abort(err) => PolicyResult {
ctx: result.ctx,
request: result.request,
output: PolicyOutput::Abort(err),
},
PolicyOutput::Retry => PolicyResult {
ctx: result.ctx,
request: result.request,
output: PolicyOutput::Retry,
},
}
}
None => PolicyResult {
ctx,
request,
output: PolicyOutput::Ready(None),
},
}
}
}

0 comments on commit f04fb97

Please sign in to comment.