opendal/layers/
concurrent_limit.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt::Debug;
19use std::pin::Pin;
20use std::sync::Arc;
21use std::task::Context;
22use std::task::Poll;
23
24use futures::Stream;
25use futures::StreamExt;
26use mea::semaphore::OwnedSemaphorePermit;
27use mea::semaphore::Semaphore;
28
29use crate::raw::*;
30use crate::*;
31
32/// Add concurrent request limit.
33///
34/// # Notes
35///
36/// Users can control how many concurrent connections could be established
37/// between OpenDAL and underlying storage services.
38///
39/// All operators wrapped by this layer will share a common semaphore. This
40/// allows you to reuse the same layer across multiple operators, ensuring
41/// that the total number of concurrent requests across the entire
42/// application does not exceed the limit.
43///
44/// # Examples
45///
46/// Add a concurrent limit layer to the operator:
47///
48/// ```no_run
49/// # use opendal::layers::ConcurrentLimitLayer;
50/// # use opendal::services;
51/// # use opendal::Operator;
52/// # use opendal::Result;
53///
54/// # fn main() -> Result<()> {
55/// let _ = Operator::new(services::Memory::default())?
56///     .layer(ConcurrentLimitLayer::new(1024))
57///     .finish();
58/// Ok(())
59/// # }
60/// ```
61///
62/// Share a concurrent limit layer between the operators:
63///
64/// ```no_run
65/// # use opendal::layers::ConcurrentLimitLayer;
66/// # use opendal::services;
67/// # use opendal::Operator;
68/// # use opendal::Result;
69///
70/// # fn main() -> Result<()> {
71/// let limit = ConcurrentLimitLayer::new(1024);
72///
73/// let _operator_a = Operator::new(services::Memory::default())?
74///     .layer(limit.clone())
75///     .finish();
76/// let _operator_b = Operator::new(services::Memory::default())?
77///     .layer(limit.clone())
78///     .finish();
79///
80/// Ok(())
81/// # }
82/// ```
83#[derive(Clone)]
84pub struct ConcurrentLimitLayer {
85    operation_semaphore: Arc<Semaphore>,
86    http_semaphore: Option<Arc<Semaphore>>,
87}
88
89impl ConcurrentLimitLayer {
90    /// Create a new ConcurrentLimitLayer will specify permits.
91    ///
92    /// This permits will applied to all operations.
93    pub fn new(permits: usize) -> Self {
94        Self {
95            operation_semaphore: Arc::new(Semaphore::new(permits)),
96            http_semaphore: None,
97        }
98    }
99
100    /// Set a concurrent limit for HTTP requests.
101    ///
102    /// This will limit the number of concurrent HTTP requests made by the
103    /// operator.
104    pub fn with_http_concurrent_limit(mut self, permits: usize) -> Self {
105        self.http_semaphore = Some(Arc::new(Semaphore::new(permits)));
106        self
107    }
108}
109
110impl<A: Access> Layer<A> for ConcurrentLimitLayer {
111    type LayeredAccess = ConcurrentLimitAccessor<A>;
112
113    fn layer(&self, inner: A) -> Self::LayeredAccess {
114        let info = inner.info();
115
116        // Update http client with metrics http fetcher.
117        info.update_http_client(|client| {
118            HttpClient::with(ConcurrentLimitHttpFetcher {
119                inner: client.into_inner(),
120                http_semaphore: self.http_semaphore.clone(),
121            })
122        });
123
124        ConcurrentLimitAccessor {
125            inner,
126            semaphore: self.operation_semaphore.clone(),
127        }
128    }
129}
130
131pub struct ConcurrentLimitHttpFetcher {
132    inner: HttpFetcher,
133    http_semaphore: Option<Arc<Semaphore>>,
134}
135
136impl HttpFetch for ConcurrentLimitHttpFetcher {
137    async fn fetch(&self, req: http::Request<Buffer>) -> Result<http::Response<HttpBody>> {
138        let Some(semaphore) = self.http_semaphore.clone() else {
139            return self.inner.fetch(req).await;
140        };
141
142        let permit = semaphore.acquire_owned(1).await;
143
144        let resp = self.inner.fetch(req).await?;
145        let (parts, body) = resp.into_parts();
146        let body = body.map_inner(|s| {
147            Box::new(ConcurrentLimitStream {
148                inner: s,
149                _permit: permit,
150            })
151        });
152        Ok(http::Response::from_parts(parts, body))
153    }
154}
155
156pub struct ConcurrentLimitStream<S> {
157    inner: S,
158    // Hold on this permit until this reader has been dropped.
159    _permit: OwnedSemaphorePermit,
160}
161
162impl<S> Stream for ConcurrentLimitStream<S>
163where
164    S: Stream<Item = Result<Buffer>> + Unpin + 'static,
165{
166    type Item = Result<Buffer>;
167
168    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
169        self.inner.poll_next_unpin(cx)
170    }
171}
172
173#[derive(Debug, Clone)]
174pub struct ConcurrentLimitAccessor<A: Access> {
175    inner: A,
176    semaphore: Arc<Semaphore>,
177}
178
179impl<A: Access> LayeredAccess for ConcurrentLimitAccessor<A> {
180    type Inner = A;
181    type Reader = ConcurrentLimitWrapper<A::Reader>;
182    type Writer = ConcurrentLimitWrapper<A::Writer>;
183    type Lister = ConcurrentLimitWrapper<A::Lister>;
184    type Deleter = ConcurrentLimitWrapper<A::Deleter>;
185
186    fn inner(&self) -> &Self::Inner {
187        &self.inner
188    }
189
190    async fn create_dir(&self, path: &str, args: OpCreateDir) -> Result<RpCreateDir> {
191        let _permit = self.semaphore.acquire(1).await;
192
193        self.inner.create_dir(path, args).await
194    }
195
196    async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
197        let permit = self.semaphore.clone().acquire_owned(1).await;
198
199        self.inner
200            .read(path, args)
201            .await
202            .map(|(rp, r)| (rp, ConcurrentLimitWrapper::new(r, permit)))
203    }
204
205    async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
206        let permit = self.semaphore.clone().acquire_owned(1).await;
207
208        self.inner
209            .write(path, args)
210            .await
211            .map(|(rp, w)| (rp, ConcurrentLimitWrapper::new(w, permit)))
212    }
213
214    async fn stat(&self, path: &str, args: OpStat) -> Result<RpStat> {
215        let _permit = self.semaphore.acquire(1).await;
216
217        self.inner.stat(path, args).await
218    }
219
220    async fn delete(&self) -> Result<(RpDelete, Self::Deleter)> {
221        let permit = self.semaphore.clone().acquire_owned(1).await;
222
223        self.inner
224            .delete()
225            .await
226            .map(|(rp, w)| (rp, ConcurrentLimitWrapper::new(w, permit)))
227    }
228
229    async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Lister)> {
230        let permit = self.semaphore.clone().acquire_owned(1).await;
231
232        self.inner
233            .list(path, args)
234            .await
235            .map(|(rp, s)| (rp, ConcurrentLimitWrapper::new(s, permit)))
236    }
237}
238
239pub struct ConcurrentLimitWrapper<R> {
240    inner: R,
241
242    // Hold on this permit until this reader has been dropped.
243    _permit: OwnedSemaphorePermit,
244}
245
246impl<R> ConcurrentLimitWrapper<R> {
247    fn new(inner: R, permit: OwnedSemaphorePermit) -> Self {
248        Self {
249            inner,
250            _permit: permit,
251        }
252    }
253}
254
255impl<R: oio::Read> oio::Read for ConcurrentLimitWrapper<R> {
256    async fn read(&mut self) -> Result<Buffer> {
257        self.inner.read().await
258    }
259}
260
261impl<R: oio::Write> oio::Write for ConcurrentLimitWrapper<R> {
262    async fn write(&mut self, bs: Buffer) -> Result<()> {
263        self.inner.write(bs).await
264    }
265
266    async fn close(&mut self) -> Result<Metadata> {
267        self.inner.close().await
268    }
269
270    async fn abort(&mut self) -> Result<()> {
271        self.inner.abort().await
272    }
273}
274
275impl<R: oio::List> oio::List for ConcurrentLimitWrapper<R> {
276    async fn next(&mut self) -> Result<Option<oio::Entry>> {
277        self.inner.next().await
278    }
279}
280
281impl<R: oio::Delete> oio::Delete for ConcurrentLimitWrapper<R> {
282    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
283        self.inner.delete(path, args).await
284    }
285
286    async fn close(&mut self) -> Result<()> {
287        self.inner.close().await
288    }
289}
opendal/layers/concurrent_limit.rs

opendal/layers/
concurrent_limit.rs
