opendal_layer_concurrent_limit/
lib.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
18#![doc = include_str!("../README.md")]
19#![cfg_attr(docsrs, feature(doc_cfg))]
20#![cfg_attr(docsrs, doc(auto_cfg))]
21#![deny(missing_docs)]
22use std::future::Future;
23use std::pin::Pin;
24use std::sync::Arc;
25use std::task::Context;
26use std::task::Poll;
27
28use futures::Stream;
29use futures::StreamExt;
30use mea::semaphore::OwnedSemaphorePermit;
31use mea::semaphore::Semaphore;
32use opendal_core::raw::*;
33use opendal_core::*;
34
35/// ConcurrentLimitSemaphore abstracts a semaphore-like concurrency primitive
36/// that yields an owned permit released on drop.
37pub trait ConcurrentLimitSemaphore: Send + Sync + Clone + Unpin + 'static {
38    /// The owned permit type associated with the semaphore. Dropping it
39    /// must release the permit back to the semaphore.
40    type Permit: Send + Sync + 'static;
41
42    /// Acquire an owned permit asynchronously.
43    fn acquire(&self) -> impl Future<Output = Self::Permit> + MaybeSend;
44}
45
46impl ConcurrentLimitSemaphore for Arc<Semaphore> {
47    type Permit = OwnedSemaphorePermit;
48
49    async fn acquire(&self) -> Self::Permit {
50        self.clone().acquire_owned(1).await
51    }
52}
53
54/// `ConcurrentLimitLayer` controls how many concurrent requests OpenDAL can send
55/// to a storage service.
56///
57/// Operators that reuse the same [`ConcurrentLimitLayer`] instance share a
58/// semaphore. This lets an application enforce one total concurrent-request
59/// limit across multiple operators.
60///
61/// # Examples
62///
63/// The following example adds a concurrent limit layer to an operator:
64///
65/// ```no_run
66/// # use opendal_core::services;
67/// # use opendal_core::Operator;
68/// # use opendal_core::Result;
69/// # use opendal_layer_concurrent_limit::ConcurrentLimitLayer;
70/// #
71/// # fn main() -> Result<()> {
72/// let _ = Operator::new(services::Memory::default())?
73///     .layer(ConcurrentLimitLayer::new(1024));
74/// # Ok(())
75/// # }
76/// ```
77///
78/// Share a concurrent limit layer between the operators:
79///
80/// ```no_run
81/// # use opendal_core::services;
82/// # use opendal_core::Operator;
83/// # use opendal_core::Result;
84/// # use opendal_layer_concurrent_limit::ConcurrentLimitLayer;
85/// #
86/// # fn main() -> Result<()> {
87/// let limit = ConcurrentLimitLayer::new(1024);
88///
89/// let _operator_a = Operator::new(services::Memory::default())?
90///     .layer(limit.clone());
91/// let _operator_b = Operator::new(services::Memory::default())?
92///     .layer(limit.clone());
93/// # Ok(())
94/// # }
95/// ```
96#[derive(Clone)]
97pub struct ConcurrentLimitLayer<S: ConcurrentLimitSemaphore = Arc<Semaphore>> {
98    operation_semaphore: S,
99    http_semaphore: Option<S>,
100}
101
102impl<S: ConcurrentLimitSemaphore> std::fmt::Debug for ConcurrentLimitLayer<S> {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        f.debug_struct("ConcurrentLimitLayer")
105            .field("has_http_limit", &self.http_semaphore.is_some())
106            .finish_non_exhaustive()
107    }
108}
109
110impl ConcurrentLimitLayer<Arc<Semaphore>> {
111    /// Create a new `ConcurrentLimitLayer` with the specified number of
112    /// permits.
113    ///
114    /// These permits will be applied to all operations.
115    pub fn new(permits: usize) -> Self {
116        Self::with_semaphore(Arc::new(Semaphore::new(permits)))
117    }
118
119    /// Set a concurrent limit for HTTP requests.
120    ///
121    /// This convenience helper constructs a new semaphore with the specified
122    /// number of permits and calls [`ConcurrentLimitLayer::with_http_semaphore`].
123    /// Use [`ConcurrentLimitLayer::with_http_semaphore`] directly when reusing
124    /// a shared semaphore.
125    pub fn with_http_concurrent_limit(self, permits: usize) -> Self {
126        self.with_http_semaphore(Arc::new(Semaphore::new(permits)))
127    }
128}
129
130impl<S: ConcurrentLimitSemaphore> ConcurrentLimitLayer<S> {
131    /// Create a layer with any ConcurrentLimitSemaphore implementation.
132    ///
133    /// ```
134    /// # use std::sync::Arc;
135    /// # use mea::semaphore::Semaphore;
136    /// # use opendal_layer_concurrent_limit::ConcurrentLimitLayer;
137    /// let semaphore = Arc::new(Semaphore::new(1024));
138    /// let _layer = ConcurrentLimitLayer::with_semaphore(semaphore);
139    /// ```
140    pub fn with_semaphore(operation_semaphore: S) -> Self {
141        Self {
142            operation_semaphore,
143            http_semaphore: None,
144        }
145    }
146
147    /// Provide a custom HTTP concurrency semaphore instance.
148    pub fn with_http_semaphore(mut self, semaphore: S) -> Self {
149        self.http_semaphore = Some(semaphore);
150        self
151    }
152}
153
154impl<S: ConcurrentLimitSemaphore> Layer for ConcurrentLimitLayer<S>
155where
156    S::Permit: Send + Sync + 'static + Unpin,
157{
158    fn apply_service(&self, inner: Servicer) -> Servicer {
159        Arc::new(self.layer(inner))
160    }
161
162    fn apply_context(&self, _srv: Servicer, inner: OperationContext) -> OperationContext {
163        // Wrap the current HTTP transport so HTTP permits are held until the
164        // response body is dropped.
165        let transport = HttpTransporter::new(ConcurrentLimitHttpTransport::<S> {
166            inner: inner.http_transport().clone(),
167            http_semaphore: self.http_semaphore.clone(),
168        });
169        inner.with_http_transport(transport)
170    }
171}
172
173impl<S: ConcurrentLimitSemaphore> ConcurrentLimitLayer<S>
174where
175    S::Permit: Send + Sync + 'static + Unpin,
176{
177    fn layer(&self, inner: Servicer) -> ConcurrentLimitService<S> {
178        ConcurrentLimitService {
179            inner,
180            semaphore: self.operation_semaphore.clone(),
181        }
182    }
183}
184
185#[doc(hidden)]
186pub struct ConcurrentLimitHttpTransport<S: ConcurrentLimitSemaphore> {
187    inner: HttpTransporter,
188    http_semaphore: Option<S>,
189}
190
191impl<S: ConcurrentLimitSemaphore> HttpTransport for ConcurrentLimitHttpTransport<S>
192where
193    S::Permit: Unpin,
194{
195    async fn fetch(&self, req: http::Request<Buffer>) -> Result<http::Response<HttpBody>> {
196        let Some(semaphore) = self.http_semaphore.clone() else {
197            return self.inner.fetch(req).await;
198        };
199
200        let permit = semaphore.acquire().await;
201
202        let resp = self.inner.fetch(req).await?;
203        let (parts, body) = resp.into_parts();
204        let body = body.map_inner(|s| {
205            Box::new(ConcurrentLimitStream::<_, S::Permit> {
206                inner: s,
207                _permit: permit,
208            })
209        });
210        Ok(http::Response::from_parts(parts, body))
211    }
212}
213
214struct ConcurrentLimitStream<S, P> {
215    inner: S,
216    // Hold this permit until the HTTP body stream is dropped.
217    _permit: P,
218}
219
220impl<S, P> Stream for ConcurrentLimitStream<S, P>
221where
222    S: Stream<Item = Result<Buffer>> + Unpin + 'static,
223    P: Unpin,
224{
225    type Item = Result<Buffer>;
226
227    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
228        // Safe due to Unpin bounds on S and P (thus on Self).
229        let this = self.get_mut();
230        this.inner.poll_next_unpin(cx)
231    }
232}
233
234#[doc(hidden)]
235#[derive(Clone)]
236pub struct ConcurrentLimitService<S: ConcurrentLimitSemaphore> {
237    inner: Servicer,
238    semaphore: S,
239}
240
241impl<S: ConcurrentLimitSemaphore> std::fmt::Debug for ConcurrentLimitService<S> {
242    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243        f.debug_struct("ConcurrentLimitService")
244            .field("inner", &self.inner)
245            .finish_non_exhaustive()
246    }
247}
248
249impl<S: ConcurrentLimitSemaphore> Service for ConcurrentLimitService<S>
250where
251    S::Permit: Send + Sync + 'static + Unpin,
252{
253    type Reader = ConcurrentLimitReader<oio::Reader, S>;
254    type Writer = ConcurrentLimitWrapper<oio::Writer, S>;
255    type Lister = ConcurrentLimitWrapper<oio::Lister, S>;
256    type Deleter = ConcurrentLimitWrapper<oio::Deleter, S>;
257    type Copier = ConcurrentLimitWrapper<oio::Copier, S>;
258
259    fn info(&self) -> ServiceInfo {
260        self.inner.info()
261    }
262
263    fn capability(&self) -> Capability {
264        self.inner.capability()
265    }
266
267    async fn create_dir(
268        &self,
269        ctx: &OperationContext,
270        path: &str,
271        args: OpCreateDir,
272    ) -> Result<RpCreateDir> {
273        let _permit = self.semaphore.acquire().await;
274        self.inner.create_dir(ctx, path, args).await
275    }
276
277    fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
278        self.inner
279            .read(ctx, path, args)
280            .map(|r| ConcurrentLimitReader::new(r, self.semaphore.clone()))
281    }
282
283    fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
284        self.inner
285            .write(ctx, path, args)
286            .map(|w| ConcurrentLimitWrapper::new(w, self.semaphore.clone()))
287    }
288
289    fn copy(
290        &self,
291        ctx: &OperationContext,
292        from: &str,
293        to: &str,
294        args: OpCopy,
295        opts: OpCopier,
296    ) -> Result<Self::Copier> {
297        self.inner
298            .copy(ctx, from, to, args, opts)
299            .map(|c| ConcurrentLimitWrapper::new(c, self.semaphore.clone()))
300    }
301
302    async fn rename(
303        &self,
304        ctx: &OperationContext,
305        from: &str,
306        to: &str,
307        args: OpRename,
308    ) -> Result<RpRename> {
309        let _permit = self.semaphore.acquire().await;
310        self.inner.rename(ctx, from, to, args).await
311    }
312
313    async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
314        let _permit = self.semaphore.acquire().await;
315        self.inner.stat(ctx, path, args).await
316    }
317
318    fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
319        self.inner
320            .delete(ctx)
321            .map(|w| ConcurrentLimitWrapper::new(w, self.semaphore.clone()))
322    }
323
324    fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
325        self.inner
326            .list(ctx, path, args)
327            .map(|s| ConcurrentLimitWrapper::new(s, self.semaphore.clone()))
328    }
329
330    async fn presign(
331        &self,
332        ctx: &OperationContext,
333        path: &str,
334        args: OpPresign,
335    ) -> Result<RpPresign> {
336        let _permit = self.semaphore.acquire().await;
337        self.inner.presign(ctx, path, args).await
338    }
339}
340
341#[doc(hidden)]
342pub struct ConcurrentLimitReader<R, S> {
343    inner: R,
344    semaphore: S,
345}
346
347impl<R, S> ConcurrentLimitReader<R, S> {
348    fn new(inner: R, semaphore: S) -> Self {
349        Self { inner, semaphore }
350    }
351}
352
353impl<R: oio::Read, S: ConcurrentLimitSemaphore> oio::Read for ConcurrentLimitReader<R, S>
354where
355    S::Permit: Send + Sync + 'static + Unpin,
356{
357    async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
358        let permit = self.semaphore.acquire().await;
359        let (rp, stream) = self.inner.open(range).await?;
360        Ok((
361            rp,
362            Box::new(ConcurrentLimitWrapper::new_with_permit(
363                stream,
364                self.semaphore.clone(),
365                permit,
366            )) as Box<dyn oio::ReadStreamDyn>,
367        ))
368    }
369
370    async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
371        let _permit = self.semaphore.acquire().await;
372        self.inner.read(range).await
373    }
374}
375
376#[doc(hidden)]
377pub struct ConcurrentLimitWrapper<R, S: ConcurrentLimitSemaphore> {
378    inner: R,
379    semaphore: S,
380    // Hold this permit until the wrapped operation body is dropped.
381    permit: Option<S::Permit>,
382}
383
384impl<R, S: ConcurrentLimitSemaphore> ConcurrentLimitWrapper<R, S> {
385    fn new(inner: R, semaphore: S) -> Self {
386        Self {
387            inner,
388            semaphore,
389            permit: None,
390        }
391    }
392
393    fn new_with_permit(inner: R, semaphore: S, permit: S::Permit) -> Self {
394        Self {
395            inner,
396            semaphore,
397            permit: Some(permit),
398        }
399    }
400
401    async fn acquire(&mut self) {
402        if self.permit.is_none() {
403            self.permit = Some(self.semaphore.acquire().await);
404        }
405    }
406}
407
408impl<R: oio::ReadStream, S: ConcurrentLimitSemaphore> oio::ReadStream
409    for ConcurrentLimitWrapper<R, S>
410where
411    S::Permit: Send + Sync + 'static + Unpin,
412{
413    async fn read(&mut self) -> Result<Buffer> {
414        self.acquire().await;
415        self.inner.read().await
416    }
417}
418
419impl<R: oio::Read, S: ConcurrentLimitSemaphore> oio::Read for ConcurrentLimitWrapper<R, S>
420where
421    S::Permit: Send + Sync + 'static + Unpin,
422{
423    async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
424        self.inner.open(range).await
425    }
426
427    async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
428        self.inner.read(range).await
429    }
430}
431
432impl<R: oio::Write, S: ConcurrentLimitSemaphore> oio::Write for ConcurrentLimitWrapper<R, S>
433where
434    S::Permit: Send + Sync + 'static + Unpin,
435{
436    async fn write(&mut self, bs: Buffer) -> Result<()> {
437        self.acquire().await;
438        self.inner.write(bs).await
439    }
440
441    async fn close(&mut self) -> Result<Metadata> {
442        self.acquire().await;
443        self.inner.close().await
444    }
445
446    async fn abort(&mut self) -> Result<()> {
447        self.acquire().await;
448        self.inner.abort().await
449    }
450}
451
452impl<R: oio::List, S: ConcurrentLimitSemaphore> oio::List for ConcurrentLimitWrapper<R, S>
453where
454    S::Permit: Send + Sync + 'static + Unpin,
455{
456    async fn next(&mut self) -> Result<Option<oio::Entry>> {
457        self.acquire().await;
458        self.inner.next().await
459    }
460}
461
462impl<R: oio::Delete, S: ConcurrentLimitSemaphore> oio::Delete for ConcurrentLimitWrapper<R, S>
463where
464    S::Permit: Send + Sync + 'static + Unpin,
465{
466    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
467        self.acquire().await;
468        self.inner.delete(path, args).await
469    }
470
471    async fn close(&mut self) -> Result<()> {
472        self.acquire().await;
473        self.inner.close().await
474    }
475}
476
477impl<C: oio::Copy, S: ConcurrentLimitSemaphore> oio::Copy for ConcurrentLimitWrapper<C, S>
478where
479    S::Permit: Send + Sync + 'static + Unpin,
480{
481    async fn next(&mut self) -> Result<Option<usize>> {
482        self.acquire().await;
483        self.inner.next().await
484    }
485
486    async fn close(&mut self) -> Result<Metadata> {
487        self.acquire().await;
488        self.inner.close().await
489    }
490
491    async fn abort(&mut self) -> Result<()> {
492        self.acquire().await;
493        self.inner.abort().await
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500    use opendal_core::Operator;
501    use opendal_core::services;
502    use std::future::pending;
503    use std::sync::Arc;
504    use std::time::Duration;
505    use tokio::time::timeout;
506
507    use futures::stream;
508    use http::Response;
509
510    #[tokio::test]
511    async fn operation_semaphore_can_be_shared() {
512        let semaphore = Arc::new(Semaphore::new(1));
513        let layer = ConcurrentLimitLayer::with_semaphore(semaphore.clone());
514
515        let permit = semaphore.clone().acquire_owned(1).await;
516
517        let op = Operator::new(services::Memory::default())
518            .expect("operator must build")
519            .layer(layer);
520
521        let blocked = timeout(Duration::from_millis(50), op.stat("any")).await;
522        assert!(
523            blocked.is_err(),
524            "operation should be limited by shared semaphore"
525        );
526
527        drop(permit);
528
529        let completed = timeout(Duration::from_millis(50), op.stat("any")).await;
530        assert!(
531            completed.is_ok(),
532            "operation should proceed once permit is released"
533        );
534    }
535
536    #[tokio::test]
537    async fn operation_semaphore_limits_copy_and_rename() {
538        #[derive(Clone, Debug)]
539        struct CopyRenameBackend {
540            info: ServiceInfo,
541            capability: Capability,
542        }
543
544        impl Service for CopyRenameBackend {
545            type Reader = ();
546            type Writer = ();
547            type Lister = ();
548            type Deleter = ();
549            type Copier = ();
550
551            fn info(&self) -> ServiceInfo {
552                self.info.clone()
553            }
554
555            fn capability(&self) -> Capability {
556                self.capability
557            }
558
559            async fn create_dir(
560                &self,
561                _: &OperationContext,
562                _: &str,
563                _: OpCreateDir,
564            ) -> Result<RpCreateDir> {
565                Err(Error::new(
566                    ErrorKind::Unsupported,
567                    "operation is not supported",
568                ))
569            }
570
571            async fn stat(&self, _: &OperationContext, _: &str, _: OpStat) -> Result<RpStat> {
572                Err(Error::new(
573                    ErrorKind::Unsupported,
574                    "operation is not supported",
575                ))
576            }
577
578            fn read(&self, _ctx: &OperationContext, _: &str, _: OpRead) -> Result<Self::Reader> {
579                Err(Error::new(
580                    ErrorKind::Unsupported,
581                    "operation is not supported",
582                ))
583            }
584
585            fn write(&self, _: &OperationContext, _: &str, _: OpWrite) -> Result<Self::Writer> {
586                Err(Error::new(
587                    ErrorKind::Unsupported,
588                    "operation is not supported",
589                ))
590            }
591
592            fn delete(&self, _ctx: &OperationContext) -> Result<Self::Deleter> {
593                Err(Error::new(
594                    ErrorKind::Unsupported,
595                    "operation is not supported",
596                ))
597            }
598
599            fn list(&self, _ctx: &OperationContext, _: &str, _: OpList) -> Result<Self::Lister> {
600                Err(Error::new(
601                    ErrorKind::Unsupported,
602                    "operation is not supported",
603                ))
604            }
605
606            fn copy(
607                &self,
608                _: &OperationContext,
609                _: &str,
610                _: &str,
611                _: OpCopy,
612                _: OpCopier,
613            ) -> Result<Self::Copier> {
614                Ok(())
615            }
616
617            async fn rename(
618                &self,
619                _: &OperationContext,
620                _: &str,
621                _: &str,
622                _: OpRename,
623            ) -> Result<RpRename> {
624                Ok(RpRename::default())
625            }
626
627            async fn presign(
628                &self,
629                _: &OperationContext,
630                _: &str,
631                _: OpPresign,
632            ) -> Result<RpPresign> {
633                Err(Error::new(
634                    ErrorKind::Unsupported,
635                    "operation is not supported",
636                ))
637            }
638        }
639
640        let semaphore = Arc::new(Semaphore::new(1));
641        let layer = ConcurrentLimitLayer::with_semaphore(semaphore.clone());
642        let capability = Capability {
643            copy: true,
644            rename: true,
645            ..Default::default()
646        };
647        let op = Operator::from_parts(
648            OperationContext::default(),
649            Arc::new(CopyRenameBackend {
650                info: ServiceInfo::with_scheme("mock"),
651                capability,
652            }),
653        )
654        .layer(layer);
655
656        let permit = semaphore.clone().acquire_owned(1).await;
657
658        let copy = timeout(Duration::from_millis(50), op.copy("from", "to")).await;
659        assert!(copy.is_err(), "copy should wait for the operation permit");
660
661        let rename = timeout(Duration::from_millis(50), op.rename("from", "to")).await;
662        assert!(
663            rename.is_err(),
664            "rename should wait for the operation permit"
665        );
666
667        drop(permit);
668
669        timeout(Duration::from_millis(50), op.copy("from", "to"))
670            .await
671            .expect("copy should proceed once permit is released")
672            .expect("copy should succeed");
673        timeout(Duration::from_millis(50), op.rename("from", "to"))
674            .await
675            .expect("rename should proceed once permit is released")
676            .expect("rename should succeed");
677    }
678
679    #[tokio::test]
680    async fn operation_semaphore_held_until_copier_dropped() {
681        #[derive(Debug)]
682        struct PendingCopier;
683
684        impl oio::Copy for PendingCopier {
685            async fn next(&mut self) -> Result<Option<usize>> {
686                pending().await
687            }
688
689            async fn close(&mut self) -> Result<Metadata> {
690                pending().await
691            }
692
693            async fn abort(&mut self) -> Result<()> {
694                Ok(())
695            }
696        }
697
698        #[derive(Clone, Debug)]
699        struct CopierBackend {
700            info: ServiceInfo,
701            capability: Capability,
702        }
703
704        impl Service for CopierBackend {
705            type Reader = ();
706            type Writer = ();
707            type Lister = ();
708            type Deleter = ();
709            type Copier = PendingCopier;
710
711            fn info(&self) -> ServiceInfo {
712                self.info.clone()
713            }
714
715            fn capability(&self) -> Capability {
716                self.capability
717            }
718
719            async fn create_dir(
720                &self,
721                _: &OperationContext,
722                _: &str,
723                _: OpCreateDir,
724            ) -> Result<RpCreateDir> {
725                Err(Error::new(
726                    ErrorKind::Unsupported,
727                    "operation is not supported",
728                ))
729            }
730
731            fn read(&self, _ctx: &OperationContext, _: &str, _: OpRead) -> Result<Self::Reader> {
732                Err(Error::new(
733                    ErrorKind::Unsupported,
734                    "operation is not supported",
735                ))
736            }
737
738            fn write(&self, _: &OperationContext, _: &str, _: OpWrite) -> Result<Self::Writer> {
739                Err(Error::new(
740                    ErrorKind::Unsupported,
741                    "operation is not supported",
742                ))
743            }
744
745            fn delete(&self, _ctx: &OperationContext) -> Result<Self::Deleter> {
746                Err(Error::new(
747                    ErrorKind::Unsupported,
748                    "operation is not supported",
749                ))
750            }
751
752            fn list(&self, _ctx: &OperationContext, _: &str, _: OpList) -> Result<Self::Lister> {
753                Err(Error::new(
754                    ErrorKind::Unsupported,
755                    "operation is not supported",
756                ))
757            }
758
759            fn copy(
760                &self,
761                _: &OperationContext,
762                _: &str,
763                _: &str,
764                _: OpCopy,
765                _: OpCopier,
766            ) -> Result<Self::Copier> {
767                Ok(PendingCopier)
768            }
769
770            async fn stat(&self, _: &OperationContext, _: &str, _: OpStat) -> Result<RpStat> {
771                Ok(RpStat::new(Metadata::new(EntryMode::FILE)))
772            }
773
774            async fn rename(
775                &self,
776                _: &OperationContext,
777                _: &str,
778                _: &str,
779                _: OpRename,
780            ) -> Result<RpRename> {
781                Err(Error::new(
782                    ErrorKind::Unsupported,
783                    "operation is not supported",
784                ))
785            }
786
787            async fn presign(
788                &self,
789                _: &OperationContext,
790                _: &str,
791                _: OpPresign,
792            ) -> Result<RpPresign> {
793                Err(Error::new(
794                    ErrorKind::Unsupported,
795                    "operation is not supported",
796                ))
797            }
798        }
799
800        let semaphore = Arc::new(Semaphore::new(1));
801        let layer = ConcurrentLimitLayer::with_semaphore(semaphore.clone());
802        let capability = Capability {
803            copy: true,
804            stat: true,
805            ..Default::default()
806        };
807        let op = Operator::from_parts(
808            OperationContext::default(),
809            Arc::new(CopierBackend {
810                info: ServiceInfo::with_scheme("mock"),
811                capability,
812            }),
813        )
814        .layer(layer);
815
816        let mut copier = timeout(Duration::from_millis(50), op.copier("from", "to"))
817            .await
818            .expect("copier setup should not block")
819            .expect("copier should be created");
820
821        let copy = timeout(Duration::from_millis(50), copier.next()).await;
822        assert!(copy.is_err(), "copy body should remain pending");
823
824        // The permit is held by the active copy body, so concurrent operations
825        // must time out until the copier is dropped.
826        let blocked = timeout(Duration::from_millis(50), op.stat("any")).await;
827        assert!(
828            blocked.is_err(),
829            "stat should wait while the copier holds the permit"
830        );
831
832        drop(copier);
833
834        timeout(Duration::from_millis(50), op.stat("any"))
835            .await
836            .expect("stat should proceed once the copier is dropped")
837            .expect("stat should succeed");
838    }
839
840    #[tokio::test]
841    async fn concurrent_chunked_read_with_http_limit() {
842        use opendal_core::raw::*;
843
844        struct EchoTransport;
845
846        impl HttpTransport for EchoTransport {
847            async fn fetch(&self, req: http::Request<Buffer>) -> Result<http::Response<HttpBody>> {
848                let data = req.into_body();
849                let len = data.len() as u64;
850                let body =
851                    HttpBody::new(Box::pin(stream::once(async move { Ok(data) })), Some(len));
852                Ok(http::Response::builder()
853                    .status(http::StatusCode::OK)
854                    .body(body)
855                    .unwrap())
856            }
857        }
858
859        #[derive(Clone, Debug)]
860        struct HttpBackend {
861            info: ServiceInfo,
862            capability: Capability,
863            content: Buffer,
864        }
865
866        /// Reader returned by this backend.
867        pub struct HttpReader {
868            backend: HttpBackend,
869            ctx: OperationContext,
870        }
871
872        impl HttpReader {
873            fn new(backend: HttpBackend, ctx: OperationContext, _: &str, _: OpRead) -> Self {
874                Self { backend, ctx }
875            }
876        }
877
878        impl oio::StreamRead for HttpReader {
879            async fn open(
880                &self,
881                range: BytesRange,
882            ) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
883                let backend = &self.backend;
884                let start = range.offset() as usize;
885                let data = match range.size() {
886                    Some(sz) => backend.content.slice(start..start + sz as usize),
887                    None => backend.content.slice(start..),
888                };
889                let req = http::Request::get("http://fake").body(data).unwrap();
890                let resp = self.ctx.http_transport().fetch(req).await?;
891                let rp = RpRead::new(Metadata::new(EntryMode::FILE).with_content_length(0));
892                let stream = resp.into_body();
893
894                Ok((rp, Box::new(stream) as Box<dyn oio::ReadStreamDyn>))
895            }
896        }
897
898        impl Service for HttpBackend {
899            type Reader = oio::StreamReader<HttpReader>;
900            type Writer = ();
901            type Lister = ();
902            type Deleter = ();
903            type Copier = ();
904
905            fn info(&self) -> ServiceInfo {
906                self.info.clone()
907            }
908
909            fn capability(&self) -> Capability {
910                self.capability
911            }
912
913            async fn create_dir(
914                &self,
915                _: &OperationContext,
916                _: &str,
917                _: OpCreateDir,
918            ) -> Result<RpCreateDir> {
919                Err(Error::new(
920                    ErrorKind::Unsupported,
921                    "operation is not supported",
922                ))
923            }
924
925            fn read(
926                &self,
927                ctx: &OperationContext,
928                path: &str,
929                args: OpRead,
930            ) -> Result<Self::Reader> {
931                Ok(oio::StreamReader::new(HttpReader::new(
932                    self.clone(),
933                    ctx.clone(),
934                    path,
935                    args,
936                )))
937            }
938
939            async fn stat(&self, _: &OperationContext, _: &str, _: OpStat) -> Result<RpStat> {
940                Ok(RpStat::new(
941                    Metadata::new(EntryMode::FILE).with_content_length(self.content.len() as u64),
942                ))
943            }
944
945            fn write(&self, _: &OperationContext, _: &str, _: OpWrite) -> Result<Self::Writer> {
946                Err(Error::new(
947                    ErrorKind::Unsupported,
948                    "operation is not supported",
949                ))
950            }
951
952            fn delete(&self, _ctx: &OperationContext) -> Result<Self::Deleter> {
953                Err(Error::new(
954                    ErrorKind::Unsupported,
955                    "operation is not supported",
956                ))
957            }
958
959            fn list(&self, _ctx: &OperationContext, _: &str, _: OpList) -> Result<Self::Lister> {
960                Err(Error::new(
961                    ErrorKind::Unsupported,
962                    "operation is not supported",
963                ))
964            }
965
966            fn copy(
967                &self,
968                _: &OperationContext,
969                _: &str,
970                _: &str,
971                _: OpCopy,
972                _: OpCopier,
973            ) -> Result<Self::Copier> {
974                Err(Error::new(
975                    ErrorKind::Unsupported,
976                    "operation is not supported",
977                ))
978            }
979
980            async fn rename(
981                &self,
982                _: &OperationContext,
983                _: &str,
984                _: &str,
985                _: OpRename,
986            ) -> Result<RpRename> {
987                Err(Error::new(
988                    ErrorKind::Unsupported,
989                    "operation is not supported",
990                ))
991            }
992
993            async fn presign(
994                &self,
995                _: &OperationContext,
996                _: &str,
997                _: OpPresign,
998            ) -> Result<RpPresign> {
999                Err(Error::new(
1000                    ErrorKind::Unsupported,
1001                    "operation is not supported",
1002                ))
1003            }
1004        }
1005
1006        let content = Buffer::from(vec![0u8; 4096]);
1007        let op = Operator::from_parts(
1008            OperationContext::default(),
1009            Arc::new(HttpBackend {
1010                info: ServiceInfo::with_scheme("mock"),
1011                capability: Capability {
1012                    read: true,
1013                    stat: true,
1014                    ..Default::default()
1015                },
1016                content: content.clone(),
1017            }),
1018        )
1019        .with_context(
1020            OperationContext::new().with_http_transport(HttpTransporter::new(EchoTransport)),
1021        )
1022        .layer(ConcurrentLimitLayer::new(1024).with_http_concurrent_limit(2));
1023
1024        // chunk=256 ⇒ 16 HTTP requests, concurrent=4, but only 2 HTTP permits.
1025        let result = timeout(Duration::from_secs(5), async {
1026            op.reader_with("test")
1027                .chunk(256)
1028                .concurrent(4)
1029                .await
1030                .expect("reader must build")
1031                .read(..)
1032                .await
1033        })
1034        .await;
1035
1036        let buf = result
1037            .expect("read must not deadlock (timeout)")
1038            .expect("read must succeed");
1039        assert_eq!(buf.to_bytes(), content.to_bytes());
1040    }
1041
1042    #[tokio::test]
1043    async fn http_semaphore_holds_until_body_dropped() {
1044        struct DummyTransport;
1045
1046        impl HttpTransport for DummyTransport {
1047            async fn fetch(&self, _req: http::Request<Buffer>) -> Result<Response<HttpBody>> {
1048                let body = HttpBody::new(stream::empty(), None);
1049                Ok(Response::builder()
1050                    .status(http::StatusCode::OK)
1051                    .body(body)
1052                    .expect("response must build"))
1053            }
1054        }
1055
1056        let semaphore = Arc::new(Semaphore::new(1));
1057        let layer = ConcurrentLimitLayer::new(1).with_http_semaphore(semaphore.clone());
1058        let fetcher = ConcurrentLimitHttpTransport::<Arc<Semaphore>> {
1059            inner: HttpTransporter::new(DummyTransport),
1060            http_semaphore: layer.http_semaphore.clone(),
1061        };
1062
1063        let request = http::Request::builder()
1064            .uri("http://example.invalid/")
1065            .body(Buffer::new())
1066            .expect("request must build");
1067        let _resp = fetcher
1068            .fetch(request)
1069            .await
1070            .expect("first fetch should succeed");
1071
1072        let request = http::Request::builder()
1073            .uri("http://example.invalid/")
1074            .body(Buffer::new())
1075            .expect("request must build");
1076        let blocked = timeout(Duration::from_millis(50), fetcher.fetch(request)).await;
1077        assert!(
1078            blocked.is_err(),
1079            "http fetch should block while the body holds the permit"
1080        );
1081    }
1082}
opendal_layer_concurrent_limit/lib.rs

opendal_layer_concurrent_limit/
lib.rs
