Skip to main content

opendal_layer_hotpath/
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::pin::Pin;
23use std::sync::Arc;
24use std::task::Context;
25use std::task::Poll;
26
27use futures::Stream;
28use futures::StreamExt;
29use opendal_core::raw::*;
30use opendal_core::*;
31
32const LABEL_CREATE_DIR: &str = "opendal.create_dir";
33const LABEL_READ: &str = "opendal.read";
34const LABEL_RENAME: &str = "opendal.rename";
35const LABEL_STAT: &str = "opendal.stat";
36const LABEL_PRESIGN: &str = "opendal.presign";
37
38const LABEL_READER_READ: &str = "opendal.reader.read";
39const LABEL_WRITER_WRITE: &str = "opendal.writer.write";
40const LABEL_WRITER_CLOSE: &str = "opendal.writer.close";
41const LABEL_WRITER_ABORT: &str = "opendal.writer.abort";
42const LABEL_LISTER_NEXT: &str = "opendal.lister.next";
43const LABEL_DELETER_DELETE: &str = "opendal.deleter.delete";
44const LABEL_DELETER_CLOSE: &str = "opendal.deleter.close";
45const LABEL_COPIER_NEXT: &str = "opendal.copier.next";
46const LABEL_COPIER_CLOSE: &str = "opendal.copier.close";
47const LABEL_COPIER_ABORT: &str = "opendal.copier.abort";
48const LABEL_HTTP_FETCH: &str = "opendal.http.fetch";
49const LABEL_HTTP_BODY_POLL: &str = "opendal.http.body.poll";
50
51/// `HotpathLayer` profiles every operation with
52/// [hotpath](https://docs.rs/hotpath/).
53///
54/// # Notes
55///
56/// When `hotpath` profiling is enabled, initialize a guard via
57/// [`hotpath::HotpathGuardBuilder`] or `#[hotpath::main]` before running
58/// operations. Otherwise, hotpath will panic on the first measurement.
59///
60/// # Examples
61///
62/// ```no_run
63/// # use opendal_core::services;
64/// # use opendal_core::Operator;
65/// # use opendal_core::Result;
66/// # use opendal_layer_hotpath::HotpathLayer;
67/// #
68/// # #[tokio::main]
69/// # async fn main() -> Result<()> {
70/// let _guard = hotpath::HotpathGuardBuilder::new("opendal").build();
71/// let op = Operator::new(services::Memory::default())?
72///     .layer(HotpathLayer::new());
73/// op.write("test", "hello").await?;
74/// # Ok(())
75/// # }
76/// ```
77#[derive(Clone, Debug, Default)]
78#[non_exhaustive]
79pub struct HotpathLayer {}
80
81impl HotpathLayer {
82    /// Create a new [`HotpathLayer`].
83    pub fn new() -> Self {
84        Self::default()
85    }
86}
87
88impl Layer for HotpathLayer {
89    fn apply_service(&self, inner: Servicer) -> Servicer {
90        Arc::new(self.layer(inner))
91    }
92
93    fn apply_context(&self, _srv: Servicer, inner: OperationContext) -> OperationContext {
94        let transport = HttpTransporter::new(HotpathHttpTransport {
95            inner: inner.http_transport().clone(),
96        });
97        inner.with_http_transport(transport)
98    }
99}
100
101impl HotpathLayer {
102    fn layer(&self, inner: Servicer) -> HotpathAccessor {
103        HotpathAccessor { inner }
104    }
105}
106
107#[doc(hidden)]
108#[derive(Debug)]
109pub struct HotpathAccessor {
110    inner: Servicer,
111}
112
113impl Service for HotpathAccessor {
114    type Reader = HotpathWrapper<oio::Reader>;
115    type Writer = HotpathWrapper<oio::Writer>;
116    type Lister = HotpathWrapper<oio::Lister>;
117    type Deleter = HotpathWrapper<oio::Deleter>;
118    type Copier = HotpathWrapper<oio::Copier>;
119
120    fn info(&self) -> ServiceInfo {
121        self.inner.info()
122    }
123
124    fn capability(&self) -> Capability {
125        self.inner.capability()
126    }
127
128    async fn create_dir(
129        &self,
130        ctx: &OperationContext,
131        path: &str,
132        args: OpCreateDir,
133    ) -> Result<RpCreateDir> {
134        hotpath::measure_async(LABEL_CREATE_DIR, self.inner.create_dir(ctx, path, args)).await
135    }
136
137    fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
138        self.inner.read(ctx, path, args).map(HotpathWrapper::new)
139    }
140
141    fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
142        self.inner.write(ctx, path, args).map(HotpathWrapper::new)
143    }
144
145    fn copy(
146        &self,
147        ctx: &OperationContext,
148        from: &str,
149        to: &str,
150        args: OpCopy,
151        opts: OpCopier,
152    ) -> Result<Self::Copier> {
153        self.inner
154            .copy(ctx, from, to, args, opts)
155            .map(HotpathWrapper::new)
156    }
157
158    async fn rename(
159        &self,
160        ctx: &OperationContext,
161        from: &str,
162        to: &str,
163        args: OpRename,
164    ) -> Result<RpRename> {
165        hotpath::measure_async(LABEL_RENAME, self.inner.rename(ctx, from, to, args)).await
166    }
167
168    async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
169        hotpath::measure_async(LABEL_STAT, self.inner.stat(ctx, path, args)).await
170    }
171
172    fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
173        self.inner.delete(ctx).map(HotpathWrapper::new)
174    }
175
176    fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
177        self.inner.list(ctx, path, args).map(HotpathWrapper::new)
178    }
179
180    async fn presign(
181        &self,
182        ctx: &OperationContext,
183        path: &str,
184        args: OpPresign,
185    ) -> Result<RpPresign> {
186        hotpath::measure_async(LABEL_PRESIGN, self.inner.presign(ctx, path, args)).await
187    }
188}
189
190#[doc(hidden)]
191pub struct HotpathWrapper<R> {
192    inner: R,
193}
194
195impl<R> HotpathWrapper<R> {
196    fn new(inner: R) -> Self {
197        Self { inner }
198    }
199}
200
201impl<R: oio::ReadStream> oio::ReadStream for HotpathWrapper<R> {
202    async fn read(&mut self) -> Result<Buffer> {
203        hotpath::measure_async(LABEL_READER_READ, self.inner.read()).await
204    }
205}
206
207impl<R: oio::Read> oio::Read for HotpathWrapper<R> {
208    async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
209        let (rp, stream) = hotpath::measure_async(LABEL_READ, self.inner.open(range)).await?;
210        Ok((
211            rp,
212            Box::new(HotpathWrapper::new(stream)) as Box<dyn oio::ReadStreamDyn>,
213        ))
214    }
215
216    async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
217        hotpath::measure_async(LABEL_READER_READ, self.inner.read(range)).await
218    }
219}
220
221impl<R: oio::Write> oio::Write for HotpathWrapper<R> {
222    async fn write(&mut self, bs: Buffer) -> Result<()> {
223        hotpath::measure_async(LABEL_WRITER_WRITE, self.inner.write(bs)).await
224    }
225
226    async fn close(&mut self) -> Result<Metadata> {
227        hotpath::measure_async(LABEL_WRITER_CLOSE, self.inner.close()).await
228    }
229
230    async fn abort(&mut self) -> Result<()> {
231        hotpath::measure_async(LABEL_WRITER_ABORT, self.inner.abort()).await
232    }
233}
234
235impl<R: oio::List> oio::List for HotpathWrapper<R> {
236    async fn next(&mut self) -> Result<Option<oio::Entry>> {
237        hotpath::measure_async(LABEL_LISTER_NEXT, self.inner.next()).await
238    }
239}
240
241impl<R: oio::Delete> oio::Delete for HotpathWrapper<R> {
242    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
243        hotpath::measure_async(LABEL_DELETER_DELETE, self.inner.delete(path, args)).await
244    }
245
246    async fn close(&mut self) -> Result<()> {
247        hotpath::measure_async(LABEL_DELETER_CLOSE, self.inner.close()).await
248    }
249}
250
251impl<C: oio::Copy> oio::Copy for HotpathWrapper<C> {
252    async fn next(&mut self) -> Result<Option<usize>> {
253        hotpath::measure_async(LABEL_COPIER_NEXT, self.inner.next()).await
254    }
255
256    async fn close(&mut self) -> Result<Metadata> {
257        hotpath::measure_async(LABEL_COPIER_CLOSE, self.inner.close()).await
258    }
259
260    async fn abort(&mut self) -> Result<()> {
261        hotpath::measure_async(LABEL_COPIER_ABORT, self.inner.abort()).await
262    }
263}
264
265struct HotpathHttpTransport {
266    inner: HttpTransporter,
267}
268
269impl HttpTransport for HotpathHttpTransport {
270    async fn fetch(&self, req: http::Request<Buffer>) -> Result<http::Response<HttpBody>> {
271        let resp = hotpath::measure_async(LABEL_HTTP_FETCH, self.inner.fetch(req)).await?;
272        let (parts, body) = resp.into_parts();
273        let body = body.map_inner(|stream| Box::new(HotpathStream { inner: stream }));
274        Ok(http::Response::from_parts(parts, body))
275    }
276}
277
278struct HotpathStream<S> {
279    inner: S,
280}
281
282impl<S> Stream for HotpathStream<S>
283where
284    S: Stream<Item = Result<Buffer>> + Unpin + 'static,
285{
286    type Item = Result<Buffer>;
287
288    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
289        let _label = LABEL_HTTP_BODY_POLL;
290        hotpath::measure_block!(_label, self.inner.poll_next_unpin(cx))
291    }
292}