betterstack_tracing/

lib.rs

//! # betterstack-tracing
//!
//! A [tracing](https://docs.rs/tracing) layer for sending logs to [Betterstack](https://betterstack.com).
//!
//! ## Features
//!
//! - Non-blocking async log sending
//! - Automatic batching (size and time-based)
//! - Span context tracking
//! - Configurable error handling
//! - Connection pooling and retries
//!
//! ## Quick Start
//!
//! ```rust,no_run
//! use betterstack_tracing::BetterstackLayer;
//! use tracing_subscriber::prelude::*;
//!
//! #[tokio::main]
//! async fn main() {
//!     let config = BetterstackLayer::builder("your-token-here")
//!         .build()
//!         .expect("failed to create config");
//!
//!     let betterstack_layer = BetterstackLayer::new(config);
//!
//!     tracing_subscriber::registry()
//!         .with(betterstack_layer)
//!         .init();
//!
//!     tracing::info!("Hello from betterstack-tracing!");
//! }
//! ```

mod config;
mod error;
mod payload;
mod sender;
mod visitor;

pub use config::{BetterstackConfig, BetterstackConfigBuilder};
pub use error::{BetterstackError, Result};

use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::Subscriber;
use tracing_subscriber::{Layer, layer::Context, registry::LookupSpan};

use sender::WorkerMessage;

/// The main tracing layer for sending logs to Betterstack
pub struct BetterstackLayer {
    config: Arc<BetterstackConfig>,
    tx: mpsc::Sender<WorkerMessage>,
    worker_handle: Option<tokio::task::JoinHandle<()>>,
}

impl BetterstackLayer {
    /// Create a new BetterstackLayer with the given configuration
    pub fn new(config: BetterstackConfig) -> Self {
        let config = Arc::new(config);
        let (tx, rx) = mpsc::channel(config.channel_capacity);

        // Spawn the background worker task
        let worker_handle = tokio::spawn(sender::worker_task(rx, config.clone()));

        Self {
            config,
            tx,
            worker_handle: Some(worker_handle),
        }
    }

    /// Create a configuration builder with the given token
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use betterstack_tracing::BetterstackLayer;
    /// use std::time::Duration;
    ///
    /// let config = BetterstackLayer::builder("your-token")
    ///     .batch_size(50)
    ///     .batch_delay(Duration::from_secs(5))
    ///     .build()
    ///     .expect("failed to create config");
    ///
    /// let layer = BetterstackLayer::new(config);
    /// ```
    pub fn builder(token: impl Into<String>) -> BetterstackConfigBuilder {
        BetterstackConfigBuilder::new(token)
    }

    /// Flush any pending logs
    ///
    /// This sends a flush message to the worker and waits for it to complete.
    /// Useful before shutdown to ensure all logs are sent.
    pub async fn flush(&self) {
        let _ = self.tx.send(WorkerMessage::Flush).await;
        // Give the worker a moment to flush
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    }
}

impl<S> Layer<S> for BetterstackLayer
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
        use visitor::FieldVisitor;

        // Create a visitor to collect fields
        let mut visitor = FieldVisitor::new();
        event.record(&mut visitor);

        // Extract message from visitor or use target as fallback
        let message = visitor
            .message()
            .unwrap_or_else(|| event.metadata().target())
            .to_string();

        // Get span context if enabled
        let mut span_context = if self.config.include_span_context {
            visitor::collect_span_context(&ctx, event, &message)
        } else {
            Default::default()
        };

        // Override message from span context with the actual event message
        span_context.message = message;

        // Build the payload
        let payload = payload::LogPayload::from_event(
            event,
            visitor.into_fields(),
            span_context,
            &self.config,
        );

        // Validate payload size to comply with Betterstack limits (1 MiB max)
        match serde_json::to_vec(&payload) {
            Ok(serialized) => {
                let size = serialized.len();

                if size > config::MAX_LOG_RECORD_SIZE {
                    // Log record exceeds 1 MiB limit - drop it
                    tracing::warn!(
                        size = size,
                        max_size = config::MAX_LOG_RECORD_SIZE,
                        "Log record exceeds Betterstack size limit and will be dropped"
                    );
                    return;
                }

                if size > config::RECOMMENDED_LOG_SIZE {
                    // Log is over recommended size but still within limits
                    tracing::debug!(
                        size = size,
                        recommended_size = config::RECOMMENDED_LOG_SIZE,
                        "Log record exceeds recommended size"
                    );
                }

                // Send to worker (non-blocking)
                // If the channel is full, we drop the log to avoid blocking
                let _ = self.tx.try_send(WorkerMessage::Log(payload));
            }
            Err(e) => {
                // Failed to serialize - log error
                tracing::error!(error = %e, "Failed to serialize log payload");
            }
        }
    }
}

impl Drop for BetterstackLayer {
    fn drop(&mut self) {
        // Send shutdown signal
        let _ = self.tx.try_send(WorkerMessage::Shutdown);

        // Wait for worker to finish with a timeout
        if let Some(handle) = self.worker_handle.take() {
            // Try to get the current runtime or create a new one for cleanup
            if let Ok(runtime_handle) = tokio::runtime::Handle::try_current() {
                runtime_handle.block_on(async {
                    let _ = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await;
                });
            } else {
                // If we're not in a tokio runtime, just abort the task
                handle.abort();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builder_creates_config() {
        let config = BetterstackLayer::builder("test-token")
            .build()
            .expect("failed to create config");

        // Just verify the config was created successfully
        assert_eq!(config.token(), "test-token");
    }
}