opendal-layer-dtrace 0.57.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Dtrace layer implementation for Apache OpenDAL.

#![cfg(target_os = "linux")]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![deny(missing_docs)]

use std::ffi::CString;
use std::fmt::Debug;
use std::fmt::Formatter;

use bytes::Buf;
use opendal_core::raw::*;
use opendal_core::*;
use probe::probe_lazy;

/// Support User Statically-Defined Tracing(aka USDT) on Linux
///
/// This layer is an experimental feature, it will be enabled by `features = ["layers-dtrace"]` in Cargo.toml.
///
/// For now we have following probes:
///
/// ### For Accessor
///
/// 1. ${operation}_start, arguments: path
///     1. create_dir
///     2. read
///     3. write
///     4. stat
///     5. delete
///     6. list
///     7. presign
///
/// 2. ${operation}_end, arguments: path
///     1. create_dir
///     2. read
///     3. write
///     4. stat
///     5. delete
///     6. list
///     7. presign
///
/// ### For Reader
///
/// 1. reader_read_start, arguments: path
/// 2. reader_read_ok, arguments: path, length
/// 3. reader_read_error, arguments: path
///
/// ### For Writer
///
/// 1. writer_write_start, arguments: path
/// 2. writer_write_ok, arguments: path, length
/// 3. writer_write_error, arguments: path
/// 4. writer_abort_start, arguments: path
/// 5. writer_abort_ok, arguments: path
/// 6. writer_abort_error, arguments: path
/// 7. writer_close_start, arguments: path
/// 8. writer_close_ok, arguments: path
/// 9. writer_close_error, arguments: path
///
/// Example:
///
/// ```no_run
/// # use opendal_core::services;
/// # use opendal_core::Operator;
/// # use opendal_core::Result;
/// # use opendal_layer_dtrace::DtraceLayer;
/// #
/// # #[tokio::main]
/// # async fn main() -> Result<()> {
/// // `Accessor` provides the low level APIs, we will use `Operator` normally.
/// let op: Operator = Operator::new(services::Memory::default().root("/tmp"))?
///     .layer(DtraceLayer::new())
///     .finish();
///
/// let path = "/tmp/test.txt";
/// for _ in 1..100000 {
///     let bs = vec![0; 64 * 1024 * 1024];
///     op.write(path, bs).await?;
///     op.read(path).await?;
/// }
/// # Ok(())
/// # }
/// ```
///
/// Then you can use `readelf -n target/debug/examples/dtrace` to see the probes:
///
/// ```text
/// Displaying notes found in: .note.stapsdt
///   Owner                Data size        Description
///   stapsdt              0x00000039       NT_STAPSDT (SystemTap probe descriptors)
///     Provider: opendal
///     Name: create_dir_start
///     Location: 0x00000000000f8f05, Base: 0x0000000000000000, Semaphore: 0x00000000003649f8
///     Arguments: -8@%rax
///   stapsdt              0x00000037       NT_STAPSDT (SystemTap probe descriptors)
///     Provider: opendal
///     Name: create_dir_end
///     Location: 0x00000000000f9284, Base: 0x0000000000000000, Semaphore: 0x00000000003649fa
///     Arguments: -8@%rax
///   stapsdt              0x0000003c       NT_STAPSDT (SystemTap probe descriptors)
///     Provider: opendal
///     Name: blocking_list_start
///     Location: 0x00000000000f9487, Base: 0x0000000000000000, Semaphore: 0x0000000000364a28
///     Arguments: -8@%rax
///   stapsdt              0x0000003a       NT_STAPSDT (SystemTap probe descriptors)
///     Provider: opendal
///     Name: blocking_list_end
///     Location: 0x00000000000f9546, Base: 0x0000000000000000, Semaphore: 0x0000000000364a2a
///     Arguments: -8@%rax
///   stapsdt              0x0000003c       NT_STAPSDT (SystemTap probe descriptors)
/// ```
#[derive(Clone, Default)]
#[non_exhaustive]
pub struct DtraceLayer {}

impl DtraceLayer {
    /// Create a new [`DtraceLayer`].
    pub fn new() -> Self {
        Self::default()
    }
}

impl<A: Access> Layer<A> for DtraceLayer {
    type LayeredAccess = DTraceAccessor<A>;
    fn layer(&self, inner: A) -> Self::LayeredAccess {
        DTraceAccessor { inner }
    }
}

#[doc(hidden)]
pub struct DTraceAccessor<A: Access> {
    inner: A,
}

impl<A: Access> Debug for DTraceAccessor<A> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DTraceAccessor")
            .field("inner", &self.inner)
            .finish_non_exhaustive()
    }
}

impl<A: Access> LayeredAccess for DTraceAccessor<A> {
    type Inner = A;
    type Reader = DtraceLayerWrapper<A::Reader>;
    type Writer = DtraceLayerWrapper<A::Writer>;
    type Lister = A::Lister;
    type Deleter = A::Deleter;
    type Copier = A::Copier;

    fn inner(&self) -> &Self::Inner {
        &self.inner
    }

    async fn create_dir(&self, path: &str, args: OpCreateDir) -> Result<RpCreateDir> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, create_dir_start, c_path.as_ptr());
        let result = self.inner.create_dir(path, args).await;
        probe_lazy!(opendal, create_dir_end, c_path.as_ptr());
        result
    }

    async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, read_start, c_path.as_ptr());
        let result = self
            .inner
            .read(path, args)
            .await
            .map(|(rp, r)| (rp, DtraceLayerWrapper::new(r, &path.to_string())));
        probe_lazy!(opendal, read_end, c_path.as_ptr());
        result
    }

    async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, write_start, c_path.as_ptr());
        let result = self
            .inner
            .write(path, args)
            .await
            .map(|(rp, r)| (rp, DtraceLayerWrapper::new(r, &path.to_string())));

        probe_lazy!(opendal, write_end, c_path.as_ptr());
        result
    }

    async fn copy(
        &self,
        from: &str,
        to: &str,
        args: OpCopy,
        opts: OpCopier,
    ) -> Result<(RpCopy, Self::Copier)> {
        let c_from = CString::new(from).unwrap();
        probe_lazy!(opendal, copy_start, c_from.as_ptr());
        let result = self.inner.copy(from, to, args, opts).await;
        probe_lazy!(opendal, copy_end, c_from.as_ptr());
        result
    }

    async fn stat(&self, path: &str, args: OpStat) -> Result<RpStat> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, stat_start, c_path.as_ptr());
        let result = self.inner.stat(path, args).await;
        probe_lazy!(opendal, stat_end, c_path.as_ptr());
        result
    }

    async fn delete(&self) -> Result<(RpDelete, Self::Deleter)> {
        self.inner.delete().await
    }

    async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Lister)> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, list_start, c_path.as_ptr());
        let result = self.inner.list(path, args).await;
        probe_lazy!(opendal, list_end, c_path.as_ptr());
        result
    }

    async fn presign(&self, path: &str, args: OpPresign) -> Result<RpPresign> {
        let c_path = CString::new(path).unwrap();
        probe_lazy!(opendal, presign_start, c_path.as_ptr());
        let result = self.inner.presign(path, args).await;
        probe_lazy!(opendal, presign_end, c_path.as_ptr());
        result
    }
}

#[doc(hidden)]
pub struct DtraceLayerWrapper<R> {
    inner: R,
    path: String,
}

impl<R> DtraceLayerWrapper<R> {
    fn new(inner: R, path: &String) -> Self {
        Self {
            inner,
            path: path.to_string(),
        }
    }
}

impl<R: oio::Read> oio::Read for DtraceLayerWrapper<R> {
    async fn read(&mut self) -> Result<Buffer> {
        let c_path = CString::new(self.path.clone()).unwrap();
        probe_lazy!(opendal, reader_read_start, c_path.as_ptr());
        match self.inner.read().await {
            Ok(bs) => {
                probe_lazy!(opendal, reader_read_ok, c_path.as_ptr(), bs.remaining());
                Ok(bs)
            }
            Err(e) => {
                probe_lazy!(opendal, reader_read_error, c_path.as_ptr());
                Err(e)
            }
        }
    }
}

impl<R: oio::Write> oio::Write for DtraceLayerWrapper<R> {
    async fn write(&mut self, bs: Buffer) -> Result<()> {
        let c_path = CString::new(self.path.clone()).unwrap();
        probe_lazy!(opendal, writer_write_start, c_path.as_ptr());
        self.inner
            .write(bs)
            .await
            .map(|_| {
                probe_lazy!(opendal, writer_write_ok, c_path.as_ptr());
            })
            .inspect_err(|_| {
                probe_lazy!(opendal, writer_write_error, c_path.as_ptr());
            })
    }

    async fn abort(&mut self) -> Result<()> {
        let c_path = CString::new(self.path.clone()).unwrap();
        probe_lazy!(opendal, writer_poll_abort_start, c_path.as_ptr());
        self.inner
            .abort()
            .await
            .map(|_| {
                probe_lazy!(opendal, writer_poll_abort_ok, c_path.as_ptr());
            })
            .inspect_err(|_| {
                probe_lazy!(opendal, writer_poll_abort_error, c_path.as_ptr());
            })
    }

    async fn close(&mut self) -> Result<Metadata> {
        let c_path = CString::new(self.path.clone()).unwrap();
        probe_lazy!(opendal, writer_close_start, c_path.as_ptr());
        self.inner
            .close()
            .await
            .inspect(|_| {
                probe_lazy!(opendal, writer_close_ok, c_path.as_ptr());
            })
            .inspect_err(|_| {
                probe_lazy!(opendal, writer_close_error, c_path.as_ptr());
            })
    }
}