Skip to main content

opendal_layer_tracing/
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::fmt::Debug;
23use std::pin::Pin;
24use std::sync::Arc;
25use std::task::Context;
26use std::task::Poll;
27
28use futures::Stream;
29use futures::StreamExt;
30use opendal_core::raw::*;
31use opendal_core::*;
32use tracing::Instrument;
33use tracing::Level;
34use tracing::Span;
35use tracing::span;
36
37/// `TracingLayer` traces every operation with
38/// [tracing](https://docs.rs/tracing/).
39///
40/// # Examples
41///
42/// ## Basic Setup
43///
44/// ```no_run
45/// # use opendal_core::services;
46/// # use opendal_core::Operator;
47/// # use opendal_core::Result;
48/// # use opendal_layer_tracing::TracingLayer;
49/// #
50/// # fn main() -> Result<()> {
51/// let _ = Operator::new(services::Memory::default())?
52///     .layer(TracingLayer::new());
53/// # Ok(())
54/// # }
55/// ```
56///
57/// ## Real usage
58///
59/// ```no_run
60/// # use anyhow::Result;
61/// # use opendal_core::services;
62/// # use opendal_core::Operator;
63/// # use opendal_layer_tracing::TracingLayer;
64/// # use tracing_subscriber::prelude::*;
65/// # use tracing_subscriber::EnvFilter;
66/// #
67/// # fn main() -> Result<()> {
68/// let opentelemetry = tracing_opentelemetry::layer();
69///
70/// tracing_subscriber::registry()
71///     .with(EnvFilter::from_default_env())
72///     .with(opentelemetry)
73///     .try_init()?;
74///
75/// {
76///     let runtime = tokio::runtime::Runtime::new()?;
77///     runtime.block_on(async {
78///         let root = tracing::span!(tracing::Level::INFO, "app_start", work_units = 2);
79///         let _enter = root.enter();
80///
81///         let _ = dotenvy::dotenv();
82///         let op = Operator::new(services::Memory::default())?
83///             .layer(TracingLayer::new());
84///
85///         op.write("test", "0".repeat(16 * 1024 * 1024).into_bytes())
86///             .await?;
87///         op.stat("test").await?;
88///         op.read("test").await?;
89///         Ok::<(), opendal_core::Error>(())
90///     })?;
91/// }
92///
93/// # Ok(())
94/// # }
95/// ```
96///
97/// # Output
98///
99/// OpenDAL is using [`tracing`](https://docs.rs/tracing/latest/tracing/) for tracing internally.
100///
101/// To enable tracing output, please init one of the subscribers that `tracing` supports.
102///
103/// For example:
104///
105/// ```no_run
106/// # use tracing::dispatcher;
107/// # use tracing::Event;
108/// # use tracing::Metadata;
109/// # use tracing::span::Attributes;
110/// # use tracing::span::Id;
111/// # use tracing::span::Record;
112/// # use tracing::subscriber::Subscriber;
113/// #
114/// # pub struct FooSubscriber;
115/// # impl Subscriber for FooSubscriber {
116/// #   fn enabled(&self, _: &Metadata) -> bool { false }
117/// #   fn new_span(&self, _: &Attributes) -> Id { Id::from_u64(0) }
118/// #   fn record(&self, _: &Id, _: &Record) {}
119/// #   fn record_follows_from(&self, _: &Id, _: &Id) {}
120/// #   fn event(&self, _: &Event) {}
121/// #   fn enter(&self, _: &Id) {}
122/// #   fn exit(&self, _: &Id) {}
123/// # }
124/// # impl FooSubscriber { fn new() -> Self { FooSubscriber } }
125///
126/// let my_subscriber = FooSubscriber::new();
127/// tracing::subscriber::set_global_default(my_subscriber).expect("setting tracing default failed");
128/// ```
129///
130/// For real-world usage, please take a look at [`tracing-opentelemetry`](https://crates.io/crates/tracing-opentelemetry).
131#[derive(Clone, Debug, Default)]
132#[non_exhaustive]
133pub struct TracingLayer {}
134
135impl TracingLayer {
136    /// Create a new [`TracingLayer`].
137    pub fn new() -> Self {
138        Self::default()
139    }
140}
141
142impl Layer for TracingLayer {
143    fn apply_service(&self, inner: Servicer) -> Servicer {
144        Arc::new(self.layer(inner))
145    }
146
147    fn apply_context(&self, _srv: Servicer, inner: OperationContext) -> OperationContext {
148        // Give outbound HTTP requests and their response bodies dedicated spans.
149        let transport = HttpTransporter::new(TracingHttpTransport {
150            inner: inner.http_transport().clone(),
151        });
152        let executor = Executor::with(TracingExecutor {
153            inner: inner.executor().clone().into_inner(),
154        });
155
156        inner.with_http_transport(transport).with_executor(executor)
157    }
158}
159
160impl TracingLayer {
161    fn layer(&self, inner: Servicer) -> TracingService {
162        TracingService { inner }
163    }
164}
165
166struct TracingHttpTransport {
167    inner: HttpTransporter,
168}
169
170impl HttpTransport for TracingHttpTransport {
171    async fn fetch(&self, req: http::Request<Buffer>) -> Result<http::Response<HttpBody>> {
172        let span = span!(Level::DEBUG, "http::fetch", ?req);
173
174        let resp = self.inner.fetch(req).instrument(span.clone()).await?;
175
176        let (parts, body) = resp.into_parts();
177        // Keep response body polling inside the same HTTP fetch span.
178        let body = body.map_inner(|s| Box::new(TracingStream { inner: s, span }));
179        Ok(http::Response::from_parts(parts, body))
180    }
181}
182
183struct TracingExecutor {
184    inner: Arc<dyn Execute>,
185}
186
187impl Execute for TracingExecutor {
188    fn execute(&self, f: BoxedStaticFuture<()>) {
189        self.inner
190            .execute(Box::pin(f.instrument(Span::current())) as BoxedStaticFuture<()>)
191    }
192
193    fn timeout(&self) -> Option<BoxedStaticFuture<()>> {
194        self.inner.timeout()
195    }
196}
197
198struct TracingStream<S> {
199    inner: S,
200    span: Span,
201}
202
203impl<S> Stream for TracingStream<S>
204where
205    S: Stream<Item = Result<Buffer>> + Unpin + 'static,
206{
207    type Item = Result<Buffer>;
208
209    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
210        let _enter = self.span.clone().entered();
211        self.inner.poll_next_unpin(cx)
212    }
213}
214
215#[doc(hidden)]
216#[derive(Debug)]
217pub struct TracingService {
218    inner: Servicer,
219}
220
221impl Service for TracingService {
222    type Reader = TracingWrapper<oio::Reader>;
223    type Writer = TracingWrapper<oio::Writer>;
224    type Lister = TracingWrapper<oio::Lister>;
225    type Deleter = TracingWrapper<oio::Deleter>;
226    type Copier = oio::Copier;
227
228    fn info(&self) -> ServiceInfo {
229        self.inner.info()
230    }
231
232    fn capability(&self) -> Capability {
233        self.inner.capability()
234    }
235
236    async fn create_dir(
237        &self,
238        ctx: &OperationContext,
239        path: &str,
240        args: OpCreateDir,
241    ) -> Result<RpCreateDir> {
242        let span = span!(Level::DEBUG, "create_dir", path, ?args);
243        self.inner
244            .create_dir(ctx, path, args)
245            .instrument(span)
246            .await
247    }
248
249    fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
250        let span = span!(Level::DEBUG, "read", path, ?args);
251        self.inner
252            .read(ctx, path, args)
253            .map(|r| TracingWrapper::new(span, r))
254    }
255
256    fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
257        let span = span!(Level::DEBUG, "write", path, ?args);
258        self.inner
259            .write(ctx, path, args)
260            .map(|r| TracingWrapper::new(span, r))
261    }
262
263    fn copy(
264        &self,
265        ctx: &OperationContext,
266        from: &str,
267        to: &str,
268        args: OpCopy,
269        opts: OpCopier,
270    ) -> Result<Self::Copier> {
271        let span = span!(Level::DEBUG, "copy", from, to, ?args, ?opts);
272        let _guard = span.enter();
273        self.inner.copy(ctx, from, to, args, opts)
274    }
275
276    async fn rename(
277        &self,
278        ctx: &OperationContext,
279        from: &str,
280        to: &str,
281        args: OpRename,
282    ) -> Result<RpRename> {
283        let span = span!(Level::DEBUG, "rename", from, to, ?args);
284        self.inner
285            .rename(ctx, from, to, args)
286            .instrument(span)
287            .await
288    }
289
290    async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
291        let span = span!(Level::DEBUG, "stat", path, ?args);
292        self.inner.stat(ctx, path, args).instrument(span).await
293    }
294
295    fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
296        let span = span!(Level::DEBUG, "delete");
297        self.inner.delete(ctx).map(|r| TracingWrapper::new(span, r))
298    }
299
300    fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
301        let span = span!(Level::DEBUG, "list", path, ?args);
302        self.inner
303            .list(ctx, path, args)
304            .map(|r| TracingWrapper::new(span, r))
305    }
306
307    async fn presign(
308        &self,
309        ctx: &OperationContext,
310        path: &str,
311        args: OpPresign,
312    ) -> Result<RpPresign> {
313        let span = span!(Level::DEBUG, "presign", path, ?args);
314        self.inner.presign(ctx, path, args).instrument(span).await
315    }
316}
317
318#[doc(hidden)]
319pub struct TracingWrapper<R> {
320    span: Span,
321    inner: R,
322}
323
324impl<R> TracingWrapper<R> {
325    fn new(span: Span, inner: R) -> Self {
326        Self { span, inner }
327    }
328}
329
330impl<R: oio::ReadStream> oio::ReadStream for TracingWrapper<R> {
331    async fn read(&mut self) -> Result<Buffer> {
332        self.inner.read().instrument(self.span.clone()).await
333    }
334}
335
336impl<R: oio::Read> oio::Read for TracingWrapper<R> {
337    async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
338        let span = span!(parent: &self.span, Level::DEBUG, "reader.open", range = %range);
339        let (rp, stream) = self.inner.open(range).instrument(span.clone()).await?;
340        Ok((
341            rp,
342            Box::new(TracingWrapper::new(span, stream)) as Box<dyn oio::ReadStreamDyn>,
343        ))
344    }
345
346    async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
347        let span = span!(parent: &self.span, Level::DEBUG, "reader.read", range = %range);
348        self.inner.read(range).instrument(span).await
349    }
350}
351
352impl<R: oio::Write> oio::Write for TracingWrapper<R> {
353    async fn write(&mut self, bs: Buffer) -> Result<()> {
354        self.inner.write(bs).instrument(self.span.clone()).await
355    }
356
357    async fn abort(&mut self) -> Result<()> {
358        self.inner.abort().instrument(self.span.clone()).await
359    }
360
361    async fn close(&mut self) -> Result<Metadata> {
362        self.inner.close().instrument(self.span.clone()).await
363    }
364}
365
366impl<R: oio::List> oio::List for TracingWrapper<R> {
367    async fn next(&mut self) -> Result<Option<oio::Entry>> {
368        self.inner.next().instrument(self.span.clone()).await
369    }
370}
371
372impl<R: oio::Delete> oio::Delete for TracingWrapper<R> {
373    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
374        self.inner
375            .delete(path, args)
376            .instrument(self.span.clone())
377            .await
378    }
379
380    async fn close(&mut self) -> Result<()> {
381        self.inner.close().instrument(self.span.clone()).await
382    }
383}