1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

// Copyright 2022 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

/// Contains layers used in this crate's subscriber for node diagnostics.
pub mod layer;

/// Contains visitors that record [`Span`](tracing::Span) field information.
pub(crate) mod visitors;

use crate::{util::Flamegrapher, Error};

use fern_logger::LoggerConfig;

use tracing_log::LogTracer;
use tracing_subscriber::{layer::Layered, prelude::*, Registry};

use std::path::{Path, PathBuf};

/// Initialises a [`tracing_log::LogTracer`] that converts any incoming [`log`] records into [`tracing`] events,
/// allowing subscribers to interact with log records.
///
/// This is necessary to perform logging through the [`log`] crate whilst also setting the global logger to a
/// [`tracing::Subscriber`] implementation.
///
/// # Errors
/// - Returns an [`Error`] if this function has failed to set the global logger.
///
/// # Notes
///  - This should only be called once. Any subsequent calls will fail, since the global logger can only be set
/// once in a program's lifespan.
///  - If the global logger has already been set (by the [`log`] crate, for example), this will fail.
pub fn collect_logs() -> Result<(), log::SetLoggerError> {
    LogTracer::init()
}

/// [`Layered`](tracing_subscriber::layer::Layered) type describing the composition of the subscriber constructed
/// by the [`trace_tools`](crate) crate.
pub type TraceSubscriber =
    Layered<Option<layer::LogLayer>, Layered<Option<layer::FlamegraphFilteredLayer>, Registry>>;

/// Builder for the [`trace_tools`](crate) subscriber.
///
/// This can be used to enable/disable [`Layer`](tracing_subscriber::Layer)s provided by this crate.
#[derive(Default)]
pub struct SubscriberBuilder {
    logger_config: Option<LoggerConfig>,
    flamegraph_stack_file: Option<PathBuf>,
}

impl SubscriberBuilder {
    /// Enables the [`LogLayer`](layer::LogLayer) for this subscriber, using the parameters provided by the
    /// given [`LoggerConfig`].
    #[must_use]
    pub fn with_log_layer(mut self, logger_config: LoggerConfig) -> Self {
        self.logger_config = Some(logger_config);
        self
    }

    /// Enables the [`FlamegraphFilteredLayer`](layer::FlamegraphFilteredLayer) for this subscriber.
    ///
    /// The given path describes the desired output location of the folded stack file that is generated by
    /// this layer during runtime. This file can then be used to produce a flamegraph by a [`Flamegrapher`]
    /// instance, or by the [`inferno`] tool.
    #[must_use]
    pub fn with_flamegraph_layer<P: AsRef<Path>>(mut self, folded_stack_file: P) -> Self {
        self.flamegraph_stack_file = Some(folded_stack_file.as_ref().to_path_buf());
        self
    }

    /// Builds and returns the [`TraceSubscriber`].
    ///
    /// # Errors
    ///  - Creation of the [`FlamegraphFilteredLayer`](layer::FlamegraphFilteredLayer) failed.
    ///  - Creation of the [`LogLayer`](layer::LogLayer) failed.
    ///
    /// # Notes
    ///  - This method calls the [`collect_logs`] function. Any [`log`] records emitted will be converting
    /// into [`tracing`] events, and therefore any external functionality that deals with [`log`] records
    /// may no longer function as expected.
    ///  - This method does *not* set the global subscriber. As such, a call to `finish` can be used to
    /// further extend the return subscriber with external [`Layer`](tracing_subscriber::Layer)s.
    pub fn finish(mut self) -> Result<(TraceSubscriber, Option<Flamegrapher>), Error> {
        let (flamegraph_layer, flamegrapher) = self.build_flamegraph_layer()?;
        let log_layer = self.build_log_layer()?;

        let subscriber = tracing_subscriber::registry()
            .with(flamegraph_layer)
            .with(log_layer);

        Ok((subscriber, flamegrapher))
    }

    /// Builds the [`TraceSubscriber`] and sets it as the global default subscriber.
    ///
    /// Returns a `Result` over an [`Option<Flamegrapher>`](Flamegrapher). The returned option is `Some` if
    /// the [`LogLayer`](layer::LogLayer) is enabled and has been successfully initialised. If the
    /// [`LogLayer](layer::LogLayer) has not been enabled with the builder, it is fine to ignore this value.
    ///
    /// # Errors
    ///  - Creation of the [`FlamegraphFilteredLayer`](layer::FlamegraphFilteredLayer) failed.
    ///  - Creation of the [`LogLayer`](layer::LogLayer) failed.
    ///
    /// # Panics
    /// This method will panic if the flamegraph layer is enabled and the program is not built with
    /// `--cfg tokio_unstable`.
    ///
    /// # Notes
    ///  - This method calls the [`collect_logs`] function. Any [`log`] records emitted will be converting
    /// into [`tracing`] events, and therefore any external functionality that deals with [`log`] records
    /// may no longer function as expected.
    ///  - This method sets the global subscriber. Any further attempts to set the global subscriber
    /// (including another call to this method) will fail.
    ///  - The subscriber initialised by this method cannot be extended.
    pub fn init(mut self) -> Result<Option<Flamegrapher>, Error> {
        let (flamegraph_layer, flamegrapher) = self.build_flamegraph_layer()?;
        let log_layer = self.build_log_layer()?;

        let subscriber = tracing_subscriber::registry()
            .with(flamegraph_layer)
            .with(log_layer);

        subscriber.init();

        Ok(flamegrapher)
    }

    fn build_log_layer(&mut self) -> Result<Option<layer::LogLayer>, Error> {
        if self.logger_config.is_some() {
            collect_logs().map_err(|err| Error::LogLayer(err.into()))?;
        }

        self.logger_config
            .take()
            .map(layer::log_layer)
            .map_or(Ok(None), |res| res.map(Some))
    }

    fn build_flamegraph_layer(
        &mut self,
    ) -> Result<(Option<layer::FlamegraphFilteredLayer>, Option<Flamegrapher>), Error> {
        self.flamegraph_stack_file
            .take()
            .map_or(Ok((None, None)), |stack_file| {
                layer::flamegraph_layer(stack_file)
                    .map(|(layer, flamegrapher)| (Some(layer), Some(flamegrapher)))
            })
    }
}

/// Returns a new, default [`SubscriberBuilder`].
pub fn build() -> SubscriberBuilder {
    SubscriberBuilder::default()
}