lance_core/utils/
tracing.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use futures::Stream;
5use pin_project::pin_project;
6use tracing::Span;
7
8#[pin_project]
9pub struct InstrumentedStream<I: Stream> {
10    #[pin]
11    stream: I,
12    span: Span,
13}
14
15impl<I: Stream> Stream for InstrumentedStream<I> {
16    type Item = I::Item;
17
18    fn poll_next(
19        self: std::pin::Pin<&mut Self>,
20        cx: &mut std::task::Context<'_>,
21    ) -> std::task::Poll<Option<Self::Item>> {
22        let this = self.project();
23        let _guard = this.span.enter();
24        this.stream.poll_next(cx)
25    }
26}
27
28// It would be nice to call the method in_current_span but sadly the Instrumented trait in
29// the tracing crate already stole the name for all Sized types
30pub trait StreamTracingExt {
31    /// All calls to poll the stream will be done in the context of the current span (when this method is called)
32    fn stream_in_current_span(self) -> InstrumentedStream<Self>
33    where
34        Self: Stream,
35        Self: Sized;
36
37    fn stream_in_span(self, span: Span) -> InstrumentedStream<Self>
38    where
39        Self: Stream,
40        Self: Sized;
41}
42
43impl<S: Stream> StreamTracingExt for S {
44    fn stream_in_current_span(self) -> InstrumentedStream<Self>
45    where
46        Self: Stream,
47        Self: Sized,
48    {
49        self.stream_in_span(Span::current())
50    }
51
52    fn stream_in_span(self, span: Span) -> InstrumentedStream<Self>
53    where
54        Self: Stream,
55        Self: Sized,
56    {
57        InstrumentedStream { stream: self, span }
58    }
59}
60
61pub const TRACE_FILE_AUDIT: &str = "lance::file_audit";
62pub const AUDIT_MODE_CREATE: &str = "create";
63pub const AUDIT_MODE_DELETE: &str = "delete";
64pub const AUDIT_MODE_DELETE_UNVERIFIED: &str = "delete_unverified";
65pub const AUDIT_TYPE_DELETION: &str = "deletion";
66pub const AUDIT_TYPE_MANIFEST: &str = "manifest";
67pub const AUDIT_TYPE_INDEX: &str = "index";
68pub const AUDIT_TYPE_DATA: &str = "data";
69pub const TRACE_FILE_CREATE: &str = "create";
70pub const TRACE_IO_EVENTS: &str = "lance::io_events";
71pub const IO_TYPE_OPEN_SCALAR: &str = "open_scalar_index";
72pub const IO_TYPE_OPEN_VECTOR: &str = "open_vector_index";
73pub const IO_TYPE_OPEN_FRAG_REUSE: &str = "open_frag_reuse_index";
74pub const IO_TYPE_OPEN_MEM_WAL: &str = "open_mem_wal_index";
75pub const IO_TYPE_LOAD_VECTOR_PART: &str = "load_vector_part";
76pub const IO_TYPE_LOAD_SCALAR_PART: &str = "load_scalar_part";
77pub const TRACE_EXECUTION: &str = "lance::execution";
78pub const EXECUTION_PLAN_RUN: &str = "plan_run";
79pub const TRACE_DATASET_EVENTS: &str = "lance::dataset_events";
80pub const DATASET_WRITING_EVENT: &str = "writing";
81pub const DATASET_COMMITTED_EVENT: &str = "committed";
82pub const DATASET_DROPPING_COLUMN_EVENT: &str = "dropping_column";
83pub const DATASET_DELETING_EVENT: &str = "deleting";
84pub const DATASET_COMPACTING_EVENT: &str = "compacting";
85pub const DATASET_CLEANING_EVENT: &str = "cleaning";
86pub const DATASET_LOADING_EVENT: &str = "loading";
lance_core/utils/tracing.rs

lance_core/utils/
tracing.rs
