Warning
Info
I'm implementing a rate limiting middleware for tower, whose Service trait requires an associated type Future for the returned future.
To check whether the request should be rate limited, my middleware makes an asynchronous call to a data store like Redis. But since async functions return anonymous futures, and impl Trait in associated types is still unstable, I don't know how to define Future.
I found another project solving the same problem with Box::pin and I'm currently doing the same thing, but I wonder if there's a way to not have to use the heap in this case.
This is my code for the service (the lifetime for K::Key is probably not ideal but I'm tyring to fix the Box problem first):
impl<S, K, ReqBody, ResBody> Service<Request<ReqBody>> for RateLimit<S, K>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>>,
K: KeyExtractor<ErrorBody = ResBody>,
for<'a> K::Key: 'a,
ResBody: Default,
{
type Response = S::Response;
type Error = S::Error;
type Future = ResponseFuture<S::Future, K::ErrorBody>;
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
let key = match self.key_extractor.extract(&req) {
Ok(key) => key,
Err(res) => return ResponseFuture::ready(res),
};
let check_future = {
let mut store = self.store.clone();
let key = key.clone();
let limit = self.limit;
let window = self.window;
Box::pin(async move { store.check_and_increment(key, limit, window).await })
};
ResponseFuture::pending(self.inner.call(req), check_future, self.fail_open)
}
}
#[pin_project(project = ResponseFutureProj)]
pub enum ResponseFuture<F, ErrorBody> {
Pending {
#[pin]
inner: F,
#[pin]
check_future: CheckFuture,
check_result: Option<RateLimitStatus>,
fail_open: bool,
},
Ready(Option<Response<ErrorBody>>),
}
and the signature for check_and_increment:
async fn check_and_increment<K>(
&mut self,
key: K,
limit: u16,
window: u32,
) -> RedisResult<RateLimitStatus>
where
K: ToRedisArgs,
{
RedisResult is just a type alias for Result<T, RedisError> (sorry can't hyperlink because as a new user on the forum I can only put 5 links max on a post).
The code is MPL-2.0. If you post some code solution but don't want me to integrate it into the rest of the project under MPL-2.0, let me know!
The rest of the code is probably not so relevant, but they are on ~liliace/tower-rate-limit-redis - A rate-limiting middleware for Tower using Redis as the data store - sourcehut git if you want more context of the code.
1 post - 1 participant
🏷️ Rust_feed