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    type Composer = oio::Composer;
228
229    fn info(&self) -> ServiceInfo {
230        self.inner.info()
231    }
232
233    fn capability(&self) -> Capability {
234        self.inner.capability()
235    }
236
237    fn compose(&self, ctx: &OperationContext, to: &str, args: OpCompose) -> Result<Self::Composer> {
238        self.inner.compose(ctx, to, args)
239    }
240
241    async fn create_dir(
242        &self,
243        ctx: &OperationContext,
244        path: &str,
245        args: OpCreateDir,
246    ) -> Result<RpCreateDir> {
247        let span = span!(Level::DEBUG, "create_dir", path, ?args);
248        self.inner
249            .create_dir(ctx, path, args)
250            .instrument(span)
251            .await
252    }
253
254    fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
255        let span = span!(Level::DEBUG, "read", path, ?args);
256        self.inner
257            .read(ctx, path, args)
258            .map(|r| TracingWrapper::new(span, r))
259    }
260
261    fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
262        let span = span!(Level::DEBUG, "write", path, ?args);
263        self.inner
264            .write(ctx, path, args)
265            .map(|r| TracingWrapper::new(span, r))
266    }
267
268    fn copy(
269        &self,
270        ctx: &OperationContext,
271        from: &str,
272        to: &str,
273        args: OpCopy,
274    ) -> Result<Self::Copier> {
275        let span = span!(Level::DEBUG, "copy", from, to, ?args);
276        let _guard = span.enter();
277        self.inner.copy(ctx, from, to, args)
278    }
279
280    async fn rename(
281        &self,
282        ctx: &OperationContext,
283        from: &str,
284        to: &str,
285        args: OpRename,
286    ) -> Result<RpRename> {
287        let span = span!(Level::DEBUG, "rename", from, to, ?args);
288        self.inner
289            .rename(ctx, from, to, args)
290            .instrument(span)
291            .await
292    }
293
294    async fn restore(
295        &self,
296        ctx: &OperationContext,
297        path: &str,
298        args: OpRestore,
299    ) -> Result<RpRestore> {
300        let span = span!(Level::DEBUG, "restore", path, ?args);
301        self.inner.restore(ctx, path, args).instrument(span).await
302    }
303
304    async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
305        let span = span!(Level::DEBUG, "stat", path, ?args);
306        self.inner.stat(ctx, path, args).instrument(span).await
307    }
308
309    fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
310        let span = span!(Level::DEBUG, "delete");
311        self.inner.delete(ctx).map(|r| TracingWrapper::new(span, r))
312    }
313
314    fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
315        let span = span!(Level::DEBUG, "list", path, ?args);
316        self.inner
317            .list(ctx, path, args)
318            .map(|r| TracingWrapper::new(span, r))
319    }
320
321    async fn presign(
322        &self,
323        ctx: &OperationContext,
324        path: &str,
325        args: OpPresign,
326    ) -> Result<RpPresign> {
327        let span = span!(Level::DEBUG, "presign", path, ?args);
328        self.inner.presign(ctx, path, args).instrument(span).await
329    }
330}
331
332#[doc(hidden)]
333pub struct TracingWrapper<R> {
334    span: Span,
335    inner: R,
336}
337
338impl<R> TracingWrapper<R> {
339    fn new(span: Span, inner: R) -> Self {
340        Self { span, inner }
341    }
342}
343
344impl<R: oio::ReadStream> oio::ReadStream for TracingWrapper<R> {
345    async fn read(&mut self) -> Result<Buffer> {
346        self.inner.read().instrument(self.span.clone()).await
347    }
348}
349
350impl<R: oio::Read> oio::Read for TracingWrapper<R> {
351    async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
352        let span = span!(parent: &self.span, Level::DEBUG, "reader.open", range = %range);
353        let (rp, stream) = self.inner.open(range).instrument(span.clone()).await?;
354        Ok((
355            rp,
356            Box::new(TracingWrapper::new(span, stream)) as Box<dyn oio::ReadStreamDyn>,
357        ))
358    }
359
360    async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
361        let span = span!(parent: &self.span, Level::DEBUG, "reader.read", range = %range);
362        self.inner.read(range).instrument(span).await
363    }
364}
365
366impl<R: oio::Write> oio::Write for TracingWrapper<R> {
367    async fn write(&mut self, bs: Buffer) -> Result<()> {
368        self.inner.write(bs).instrument(self.span.clone()).await
369    }
370
371    async fn copy_from(&mut self, path: &str, args: OpRead, range: BytesRange) -> Result<()> {
372        self.inner
373            .copy_from(path, args, range)
374            .instrument(self.span.clone())
375            .await
376    }
377
378    async fn abort(&mut self) -> Result<()> {
379        self.inner.abort().instrument(self.span.clone()).await
380    }
381
382    async fn close(&mut self) -> Result<Metadata> {
383        self.inner.close().instrument(self.span.clone()).await
384    }
385}
386
387impl<R: oio::List> oio::List for TracingWrapper<R> {
388    async fn next(&mut self) -> Result<Option<oio::Entry>> {
389        self.inner.next().instrument(self.span.clone()).await
390    }
391}
392
393impl<R: oio::Delete> oio::Delete for TracingWrapper<R> {
394    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
395        self.inner
396            .delete(path, args)
397            .instrument(self.span.clone())
398            .await
399    }
400
401    async fn close(&mut self) -> Result<()> {
402        self.inner.close().instrument(self.span.clone()).await
403    }
404}